Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs
Anonymous Author(s)
AffiliationAddress email
Abstract
Despite recent success in using the invariance principle for out-of-distribution(OOD) generalization on Euclidean data (e.g., images), studies on graph data arestill limited.
Different from images, the complex nature of graphs poses uniquechallenges to adopting the invariance principle.
In particular, distribution shifts ongraphs can appear in a variety of forms such as attributes and structures, making itdifficult to identify the invariance.
Moreover, domain or environment partitions,which are often required by OOD methods on Euclidean data, could be highlyexpensive to obtain for graphs.
To bridge this gap, we propose a new framework,called C ausality
Inspired I nvariant G raph Le A rning (CIGA), to capture the invariance of graphs for guaranteed OOD generalization under various distribution shifts.
Specifically, we characterize potential distribution shifts on graphs with causalmodels, concluding that OOD generalization on graphs is achievable when modelsfocus only on subgraphs containing the most information about the causes of labels.
Accordingly, we propose an information-theoretic objective to extract the desiredsubgraphs that maximally preserve the invariant intra-class information.
Learningwith these subgraphs is immune to distribution shifts.
Extensive experiments onboth synthetic and real-world datasets, including a challenging setting in AI-aideddrug discovery, validate the superior OOD generalization ability of CIGA.
Introduction
Graph representation learning with graph neural networks (GNNs) has gained great success in tasksinvolving relational information [45, 35, 98, 105, 106].
However, it assumes that the training andtest graphs are drawn from the same distribution, which is often violated in reality [37, 47, 38, 40].
The mismatch between training and test distributions, i.e., distribution shifts , introduced by someunderlying environmental factors related to data collection or processing, could seriously degradethe performance of deployed models [7, 24].
Such out-of-distribution (OOD) generalization failuresbecome the major roadblock for practical applications of graph representation learning [40].
Meanwhile, enabling OOD generalization on regular Euclidean data has received surging attentionand several solutions were proposed [4, 80, 10, 49, 23, 48, 2].
In particular, the invariance principlefrom causality is at the heart of those works [75, 73, 78].
The principle leverages the Independent
Causal Mechanism (ICM) assumption [73, 76] and implies that, model predictions that only focus onthe causes of the label can stay invariant to a large class of distribution shifts [75, 4, 2].
Despite the success of the invariance principle on Euclidean data, the complex nature of graphs raisesseveral new challenges that prohibit direct adoptions of the principle.
First, distribution shifts ongraphs are more complicated.
They can happen at both attribute-level and structure-level, and beobserved in multiple forms such as graph sizes, subgraph densities and homophily [112, 11, 101].
On the other hand, each of the shifts can spuriously correlate with labels in different modes [4, 70, 2].
“Cycle” 𝑔 “House”
Classification Invariant Subgraph Identification ෡𝑮 𝒄
Cycle
House (a)
Algorithm
OODGuarantee Regime E Known SCMSupport
IRM[4]
Yes
R
Yes
PIIF IB-IRM[2]
Yes
R
Yes
PIIF&FIIF EIIL[23]
Yes
R
No
PIIF DANN[31]
N/A
R
Yes
N/A MatchDG[60]
N/A
R
Yes
FIIF GroupDro[80]
N/A
R
Yes
N/A CNC[122]
N/A
R
No
N/A GIB[119]
Yes
G
No
FIIF DIR[103]
No
G
No
FIIF CIGA (Ours)
Yes
G
No
PIIF&FIIF (b)
Consequently, the entangled complex distribution shifts make it more difficult to identify and capturethe invariance on graphs.
Second, OOD algorithms developed and analyzed on Euclidean dataoften require additional environment (or domain) labels for distinguishing the sources of distributionshifts [4].
However, the environment labels could be highly expensive to obtain and thus oftenunavailable for graphs, as collecting the labels usually requires expert knowledge due to the abstractionof graphs [37].
These challenges render the problem studied in this paper even more challenging:
How could one generalize the invariance principle to enable OOD generalization on graphs?
To solve the above problem, we propose C ausality
Inspired I nvariant G raph Le A rning (CIGA), anew framework for capturing the invariance of graphs to enable guaranteed OOD generalization underdifferent distribution shifts.
Specifically, we build three Structural Causal Models (SCMs) [73] tocharacterize the distribution shifts that could happen on graphs: one is to model the graph generationprocess, and the other two are to model two possible interactions between invariant and spuriousfeatures during the graph generation, i.e., Fully Informative Invariant Feature (FIIF) and
Partially
Informative Invariant Feature (PIIF) (Sec. 2.2).
Then, we generalize the invariance principle tographs for OOD generalization: GNN models are invariant to distribution shifts if they focus onlyon an invariant and critical subgraph G c that contains the most of the information in G about theunderlying causes of the label.
Thus, the problem of achieving OOD generalization on graphs can berephrased into two processes: invariant subgraph identification and label prediction.
Accordingly,shown as Fig. 1(a), we introduce a prototypical invariant graph learning algorithm that decomposesa GNN into: a) a featurizer g for identifying the underlying invariant subgraph G c from G ; b) aclassifier f c for making predictions based on G c .
To extract the desired subgraph G c , we derive aninformation-theoretic objective for the featurizer to identify subgraphs that maximally preserves theinvariant intra-class information across a set of different (unknown) environments.
We theoreticallyshow that this approach can provably identify the underlying G c under mild assumptions (Sec. 3).
Experiments on 16 synthetic and real-world datasets with various distribution shifts, including achallenging setting from AI-aided drug discovery [40], show that CIGA can significantly outperformall of existing methods up to 10% , demonstrating its promising OOD generalization ability (Sec. 4).
Related Work.
We review existing methods that might improve the OOD generalization on graphs,summarize the main differences between our solution and them in Table 1(b), and leave thoroughdiscussions to Appendix B.2.
On Euclidean data, Invariant Learning [4, 23, 2], Group Distributionally
Robust Optimization [49, 80, 122], Domain Adaption and Domain Generalization [31, 92, 52, 27, 60,99] are three widely adopted approaches to enable OOD generalization.
However, they all have theirown limitations when being applied to graphs.
First, previous invariant learning methods are mostlydeveloped and analyzed for Euclidean data [4, 2, 23], or under specific SCM assumptions [4], makingthe theoretical results hardly able to generalize to the complicated graph data [79] that can havemultiple types of distribution shifts [70].
Group Distributionally Robust Optimization that minimizesthe gap between worst group risk and average risk [49, 80, 122], and Domain Adaption/Generalizationmethods that aim to learn class-conditional domain invariant representations [31, 92, 52, 27, 99], cannot guarantee a min-max optimal predictor without additional assumptions [124, 4, 2].
Moreover,most existing methods require environment labels that are however expensive to obtain in graphs,which limits their applications to graphs [4, 49, 2, 80, 31, 92, 27, 60].
In contrast, we aim to develop
OOD algorithms for graphs that are provably generalizable under different types of distribution shifts.
Another line of relevant works is about GNN explainability that aims to find a subgraph of the inputas the explanation for a GNN prediction [115, 120].
Although some may leverage causality to justifythe generated explanation [53], they mostly focus on understanding the predictions of GNNs insteadof for OOD generalization.
The most close works to ours are two interpretable GNNs that aim toexplicitly extract a subgraph for both predictions and explanations guided by information theory [119]and causality [103], respectively.
However, they focus on graphs and shifts generated under a specific
SCM.
Although one of them can provide theoretical guarantee for OOD generalization [119] by usingthe information bottleneck criteria [2], they would inevitably fail to generalize to graphs generatedunder different SCMs.
More discussions about the failure are given in Appendix D.4.
Besides,
Bevilacqua et al.
[11] also discuss OOD generalization on graphs but limited a specific graph familyand graph size shifts.
Wu et al.
[102] propose OOD generalization algorithms on graphs for the taskof node classification, also limited to graphs and shifts under a specific SCM.
To the best of our knowledge, there is no existing work that could handle more comprehensive graphdistribution shifts than CIGA, while also achieving provable OOD generalization performance.
OOD Generalization on Graphs through the Lens of Causality2.1 Problem Setup
In this work, we focus on OOD generalization in graph classification.
Specifically, we are given a setof graph datasets D
= {D e } e collected from multiple environments E all .
Samples ( G ei , Y ei ) ∈ D efrom the same environment are considered as drawn independently from an identical distribution P e .
A
GNN ρ ◦ h generically has an encoder h : G → R h that learns a meaningful representation h G for eachgraph G to help predict the label ˆ Y G
= ρ ( h G ) with a downstream classifier ρ : R h
→ Y .
The goal of
OOD generalization on graphs is to train a GNN ρ ◦ h with data from training environments D tr ={D e } e ∈E tr ⊆E all that generalizes well to all (unseen) environments, i.e., to minimize max e ∈E all R e ,where R e is the empirical risk of ρ ◦ h under environment e [96, 4].
We leave more details about thebackground of GNN for graph classification and invariant learning in Appendix B.1.
It is known that OOD generalization is impossible without assumptions on the environments E all [73,2].
Thus, we will first formulate the data generation process with structural causal model and latentvariable model [73, 76, 50], to characterize the distribution shifts that could happen on graphs.
Then,we investigate whether the existing methods are generalizable under these distribution shifts.
2.2 Graph Generation Process
[Figure]
We take a latent-variable model perspective on the graph generation process and assume that the graph is [Figure] generated through a mapping f gen :
Z → G , where Z ⊆ R n is the latent space and G = ∪ ∞ N =1 { 0 , 1 } N ×
R N × d is the graph space.
Let E denote environments.
Following previous works [50, 2], we partition the latent variable from Z into an invariantpart C ∈ C =
R n c and a varying part
S ∈ S =
R n s , s.t., n = n c + n s , according to whether they are affected by E .
Similarly in images,
C and S can represent content and style while E can refer to the locations where the images aretaken [7, 123, 50].
Furthermore, C and S control the generation of the observed graphs (Assumption 2.1) and can have multiple types of interactions at the latent space (Assumptions 2.2, 2.3).
Graph generation model.
We elaborate the SCM for the graph generation process in Assumption 2.1and Fig. 2(a), where noises in the structural equations are omitted for simplicity [76].
Assumption 2.1 (Graph Generation Structural Causal Model) .
[Equation]
In Assumption 2.1, f gen is decomposed into f G c gen , f G s gen and f G gen to control the generation of G c , G s ,and G , respectively.
Among them, G c inherits the invariant information of C that would not beaffected by the interventions (or changes) of E [73, 76].
For example, certain properties of a moleculecan usually be described by a sub-molecule, or a functional group, which is invariant across differentspecies or assays [12, 91, 40].
On the contrary, the generation of G s and G will be affected by theenvironment E through S .
Thus, graphs collected from different environments (or domains) can havedifferent distributions of structure-level properties (e.g., graph sizes [11, 101]) as well as featurelevel properties (e.g., homophily [61, 17]).
In fact, Assumption 2.1 is compatible with many graphgeneration models by specifying the function classes of f G c gen , f G s gen and f G gen [88, 56, 116, 58].
Sinceour goal is to characterize the potential distribution shifts in Assumption 2.1, we focus on building ageneral SCM that is compatible to many graph families and leave graph family specifications and theirimplications to OOD generalization in future works. More discussions are provided in Appendix C.
Interactions at latent space.
Following previous works [4, 2], we categorize the latent interactionsbetween C and S into Fully Informative Invariant Features (FIIF, Fig. 2(b)) and Partially Informative
Invariant Features (PIIF, Fig. 2(c)) 1 , depending on whether the latent invariant part C is fullyinformative about label Y , i.e., ( S, E ) ⊥⊥ Y | C .
Formal definitions of the corresponding SCMs aregiven as follows, where noises are omitted for simplicity [73, 76].
Assumption 2.2 (FIIF Structural Causal Model) .
Y :
= f inv ( C ) , S :
= f spu ( C, E ) , G :=
f gen ( C, S ) .
Assumption 2.3 (PIIF Structural Causal Model) .
Y :
= f inv ( C ) , S :
= f spu ( Y, E ) , G :=
f gen ( C, S ) .
In the two SCMs above, f gen corresponds to the graph generation process in Assumption 2.1, and f spuis the mechanism describing how S is affected by C and E at the latent space.
By definition, S isdirectly controlled by C in FIIF and indirectly controlled by C through Y in PIIF, which can exhibitdifferent behaviors in the observed distribution shifts.
In practice, performances of OOD algorithmscan degrade dramatically if one of FIIF or PIIF is excluded [5, 70].
This issue can be more seriousin graphs, since different distribution shifts can have different interaction modes at the latent space.
Moreover, f inv : C → Y indicates the labelling process, which assigns labels Y for the corresponding
G merely based on C .
Consequently, C is better clustered than S when given Y [13, 15, 85, 86],which also serves as the necessary separation assumption for a classification task [68, 16, 64].
Assumption 2.4 (Better Clustered Invariant Features) .
H ( C | Y ) ≤ H ( S | Y ) .
2.3 Challenges of OOD Generalization on Graphs
Built upon the graph generation process, we can formally derive the desired GNN that is able togeneralize to OOD graphs under different distribution shifts, which implies the invariant GNN below 2 .
Definition 2.5 (Invariant GNN) .
Given a set of graph datasets {D e } e and environments E all thatfollow the same graph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutationinvariant graph encoder h : G → R h and a downstream classifier ρ : R h → Y , ρ ◦ h is an invariant
GNN if it minimizes the worst case risk among all environments, i.e., min max e ∈E all R e .
Can existing methods produce a desired invariant GNN model?
We find the answers to be negativeunfortunately.
Based on the synthetic BAMotif graph classification task [57, 103] shown in Fig.
3,we theoretically and empirically analyze whether existing methods could produce an invariant GNN,through the investigation of the following aspects.
More details and results are given in Appendix D.
Can GNNs trained with ERM generalize to OOD graphs?
As shown in Fig.
3, we find that GNNstrained with the standard empirical risk minimization (ERM) algorithm [96] are not able to generalizeto OOD graphs.
As the data biases grows stronger, the performances of GNNs drop dramatically.
Furthermore, when graph size shifts are mixed in the data, GNNs can have larger variance at low databiases, indicating the instability of learning the desired relationships for the task.
The reason is that
Training Data
Testing Data “House” “House” “Cycle”
[Figure]
0.9 Data Biases
A cc u r a c y
Struc-ERMMixed-ERMStruc-IRMMixed-IRM
[Figure]
shifts.
0.9 Data Biases
A cc u r a c y
Struc-ERMMixed-ERMStruc-IRMMixed-IRM
[Figure]
Can OOD objectives improve OOD generalization of GNNs?
Meanwhile, as shown in
Fig.
3, OOD objectives primarily developed on Euclidean data such as invariant risk minimization (IRM) [4] also cannot alleviate the problem.
On the contrary, IRM can fail catastrophically at non-linear regime if without sufficient support overlap for the test environments, i.e.,∪ e ∈E te supp ( P e ) ̸⊆ ∪ e ∈E tr supp ( P e ) [79].
In addition to IRM, the failure would also happen foralternative objectives [49, 9, 2] as proved by Rosenfeld et al.
Besides, different distributionshifts on graphs can be nested with each other where each one can have distinct spurious correlationtype, e.g., FIIF or PIIF.
OOD objectives will also fail seriously if either of the correlation types is notsupported [5, 70].
Moreover, non-trivial environment partitions or labels are required for performanceguarantee of these OOD objectives [4, 49, 80, 2].
However, collecting meaningful environmentpartitions of graphs requires expert knowledge about graph data.
Thus, the environment labels can beexpensive to obtain and are usually not available [66, 28, 37].
Alternative options such as randompartitions tend not to alleviate the issue [23, 54], as it can be trivially deemed as mini-batching.
Challenges of OOD generalization on graphs.
The aforementioned failure analysis reveals thatexisting methods or objectives fail to elicit an invariant GNN primarily due to the following twochallenges: a) Distribution shifts on graphs are more complicated where different types of spuriouscorrelations can be entangled via different graph properties; b) Environment labels are usually notavailable due to the abstraction of graphs.
Despite these challenges, we are still highly motivatedto address the following research question: Would it be possible to learn an invariant GNN that isgeneralizable under various distribution shifts by lifting the invariance principle to the graph data?
Invariance Principle for OOD Generalization on Graphs
We provide affirmative answers to the previous question by proposing a new framework, CIGA:
C ausality Inspired I nvariant G raph Le A rning.
Specifically, built upon the SCMs in Sec. 2.2, wegeneralize the invariance principle to graphs and instantiate the principle with theoretical guarantees.
3.1 Invariance for OOD Generalization on Graphs
Towards extending the invariance principle to graphs under SCMs in Sec. 2.2, we need to identify a setof variables that have stable causal relationship with Y under both FIIF and PIIF (Assumption 2.2, 2.3).
According to the ICM assumption [76], the labeling process C → Y is not informed nor influencedby other processes, implying that the conditional distribution P ( Y | C ) remains invariant to theinterventions on the environment latent variable E [73].
Consequently, for a GNN with a permutation invariant encoder h : G → R h and a downstream classifier ρ : R h
→ Y , if h can recover theinformation of C from G in the learned graph representations, then the learning of ρ resemblestraditional ERM [96] and can achieve the desired min-max optimality required by an invariant GNN(Def. 2.5).
However, recovering C from G is particularly difficult, since the generation of G from Cinvolves two causal mechanisms f G c gen and f G gen in Assumption 2.1.
The unavailability of E furtheradds up the difficulty of enforcing the independence between the learned representations and E .
3.2 Invariant Graph Learning Framework
Causal algorithmic alignment.
To enable a GNN to learn to extract the information about C from
G , we propose the CIGA framework that explicitly aligns with the two causal mechanisms f G c gen andf G gen in Assumption 2.1.
The idea of alignment in CIGA is motivated by the algorithmic reasoningresults that a neural network can learn a reasoning process better if its computation structure alignswith the process better [107, 109].
Specifically, we realize the alignment by decomposing a GNNinto two sub-components 3 : a) a featurizer GNN g : G → G c aiming to identify the desired G c ; b) aclassifier GNN f c : G c → Y that predicts the label Y based on the estimated G c , where G c refers tothe space of subgraphs of G .
Formally, the learning objectives of f c and g can be formulated as:
[Equation]
[Equation] where maximizing I ( (cid:98) G c ; Y ) is equivalent to minimizing a variational upper bound of R ( f c ( (cid:98) G c )) [3,119] that takes (cid:98) G c as inputs to predict label Y for G through f c and g , and (cid:98) G c is the estimatedsubgraph containing the information about C and hence needs to be independent of E .
Moreover, theextracted G c can either shares the same graph space with input G or has its own space with latentnode and edge features, depending on the specific graph generation process.
In practice, architecturesfrom the literature of interpretable GNNs are compatible with CIGA [120], hence can serve aspractical choices for the implementation of CIGA.
More details are given in Appendix F.
Although we can technically align with the two causal mechanisms with g and f c , trivially optimizingthis architecture cannot satisfy (cid:98) G c ⊥⊥ E .
Formally, merely maximizing I ( (cid:98) G c ; Y ) may include asubgraph from G s in (cid:98) G c since G s also shares certain mutual information with Y .
Moreover, theunavailability of E prevents the direct usage of E in enforcing the independence that is often adoptedby previous methods [4, 49, 80, 31, 92], making the identification of G c more challenging.
Optimization objective.
To mitigate this issue, we need to find and translate other properties of G cinto some differentiable and equivalent objectives to satisfy the independence constraint (cid:98) G c ⊥⊥ E . The goal of the desired objective.
We begin by considering a simplistic setting where all the invariantsubgraphs G c have the same size s c , i.e., | G c | =
s c 4 .
When maximizing I ( (cid:98) G c ; Y ) in Eq.
1, both FIIFand PIIF can introduce part of G s into (cid:98) G c .
In FIIF (Fig. 2(b)), as G c already contains the maximalpossible information in G about Y , G c is a solution to max I ( (cid:98) G c ; Y ) .
However, some subgraph of G c can be replaced by some subgraph of G s that is equally informative about Y .
In PIIF (Fig.
2(c)),there also exists some subgraph of G s that contains additional information about Y than G c , hence(cid:98) G c is more likely to involve some subgraph of G s .
Thus, the new objective needs to eliminate theauxiliary subgraphs of (cid:98) G c from G s such that the estimated (cid:98) G c can only contain G c . An important property of G c .
Under both FIIF and PIIF SCMs (Fig. 2), for G e 1 c , G e 2 c that relateto the same causal factor c under two environments e 1 and e 2 , the desired (cid:98) G e 1 c , (cid:98) G e 2 c in e 1 and e 2tend to have high mutual information, i.e., ( G e 1 c , G e 2 c ) ∈ arg max I ( (cid:98) G e 1 c ; (cid:98) G e 2 c ) .
While for G e 1 cand another G e 1 c ′ corresponding to a different c ′ ̸
= c , under the same environment e 1 , including anysubgraph from G e 1 s in (cid:98) G e 1 c , (cid:98) G e 1 c ′ will enlarge their mutual information, or in other words, ( G e 1 c , G e 1 c ′ ) ∈arg
min I ( (cid:98) G e 1 c ; (cid:98) G e 1 c ′ ) .
Thus, we can derive an important property of G c , that is, ∀ e 1 , e 2 ∈ E all ,
[Equation]
where (cid:98) G e 1 c and (cid:98) G e 2 c are the estimated invariant subgraphs corresponding to the same causal factor cunder environment e 1 and e 2 , respectively, while (cid:98) G e 2 c ′ corresponds to a different causal factor c ′ . edges or attributes. Deriving CIGA v1 based on the identified property of G c .
In practice, C is not given.
Nevertheless,since C and Y shares a stable causal relationship in both FIIF and PIIF SCMs, Y can serve as aproxy of C in Eq. 2.
Moreover, as Eq. 2 holds for any ∀ e 1 , e 2 ∈ E all , the environment superscriptscan be eliminated without affecting Eq. 2.
Furthermore, when both I ( (cid:98) G e 1 c ; (cid:98) G e 2 c | C = c ) and I ( (cid:98) G c ; Y )are maximized, I ( (cid:98) G e 1 c ; (cid:98) G e 1 c ′ | C
= c ′ , c ′ ̸
= c ) is automatically minimized, otherwise all classes willcollapse to trivial solutions which is contradictory given I ( (cid:98) G c ; Y ) being maximized.
Therefore, wecan derive an alternative objective to Eq. 1 by leveraging Eq. 2 to replace the independence condition:(
[Equation]
where (cid:101) G c = g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G is sampled from training graphs that share the samelabel Y as G .
In Theorem 3.1, we show how Eq. 3 is equivalent to Eq.
Nevertheless, Eq. 3 requiresa strong assumption on the size of G c .
However, the size of G c is usually unknown or changes fordifferent C s. In this circumstance, maximizing Eq.
2 without additional constraints will lead to thepresence of part of G s in (cid:98) G c .
For instance, (cid:98) G c =
G is a trivial solution to Eq.
3 when s c = ∞ .
Deriving CIGA v2 by resolving size constraint on G c in CIGA v1. To this end, we further resort tothe properties of G s .
In both FIIF and PIIF SCMs (Fig. 2), G s and G c can share certain overlappedinformation about Y .
When maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) and I ( (cid:98) G c ; Y ) , the appearance of partial G sin (cid:98) G c will not affect the optimality.
However, it can reduce the mutual information between theleft part (cid:98) G s =
G − (cid:98) G c and Y , i.e., I ( (cid:98) G s ; Y ) .
Therefore, by maximizing I ( (cid:98) G s ; Y ) , we can reduceincluding part of G s into (cid:98) G c .
Meanwhile, to avoid trivial solution that G c ⊆ (cid:98) G s during maximizing
I ( (cid:98) G s ; Y ) , we can leverage the better clustering property of G c implied by Assumption 2.4 to derivethe constraint
I ( (cid:98) G s ; Y ) ≤
I ( (cid:98) G c ; Y ) .
Thus, we can obtain a new objective CIGAv2 as follows:
[Equation]
where (cid:98) G c
= g ( G ) , (cid:101) G c = g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G is sampled from training graphs that sharethe same label Y as G .
We also prove the equivalence between Eq.
4 and Eq. 1 in Theorem 3.1.
3.3 Theoretical Analysis and Practical Discussions
Theorem 3.1 (CIGA Induces Invariant GNNs) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c genin Assumption 2.1 are invertible, (b) samples from each training environment are equally distributed,i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr , then:
(i).
If ∀ G c , | G c | = s c , then each solution to Eq. 3, elicits an invariant GNN (Def. 2.5).(ii).
Each solution to Eq. 4, elicits an invariant GNN (Def. 2.5).
We prove Theorem 3.1 (i) and (ii) in Appendix E.2, E.3, respectively.
Practical implementations of CIGA objectives.
After showing the power of CIGA, we introducethe practical implementations of CIGAv1 and CIGAv2 objectives.
Specifically, an exact estimate ofthe second term I ( (cid:98) G c ; (cid:101) G c | Y ) could be highly expensive [95, 8].
However, contrastive learning withsupervised sampling provides a practical solution for the approximation [42, 20, 81, 95, 8]:
[Equation]
where positive samples ( (cid:98) G c , (cid:101) G c ) are the extracted subgraphs of graphs that share the same label as
G , negative samples are those having different labels, P g ( G |Y =
Y ) is the push-forward distributionof P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label Y , 285 P ( G |Y
̸ = Y ) refers to the distribution of G given the label that is different from Y , h (cid:98) G c , h (cid:101) G c , h G icare the graph presentations of the estimated subgraphs, and ϕ is the similarity metric for graphpresentations.
As M → ∞ , Eq.
5 approximates I ( (cid:98) G c ; (cid:101) G c | Y ) , which can be regarded as a nonparameteric resubstitution entropy estimator via the von Mises-Fisher kernel density [1, 41, 100].
Thus, plugging it into Eq. 3 and Eq.
4 can relieve the issue of approximating I ( (cid:98) G c ; (cid:101) G c | Y ) in practice.
[Table]
For the implementation of I ( (cid:98) G s ; Y ) and the constraint I ( (cid:98) G s ; Y ) ≤
I ( (cid:98) G c ; Y ) in CIGAv2, a practicalchoice is to follow the idea of hinge loss, I ( (cid:98) G s ; Y )
≈ 1 N R (cid:98) G s · I ( R (cid:98) G s ≤ R (cid:98) G c ) , where N is thenumber of samples, I is an indicator function that outputs 1 when the inner condition is satisfiedotherwise 0 , and R (cid:98) G s and R (cid:98) G c are the empirical risk vector of the predictions for each sample basedon the corresponding (cid:98) G s and (cid:98) G c .
More implementation details can be found in Appendix F.
Discussions and implications of CIGA.
Although using contrastive learning to improve OODgeneralization is not new in the literature [27, 60, 122], previous methods cannot yield OOD guarantees in graph circumstances due to the highly non-linearity and the unavailability of domain labels
E .
In particular, CIGA can be reduced to directly applying contrastive learning when withoutthe decomposition for causal algorithmic alignment.
However, in the experiments we found thatmerely using the contrastive objective, i.e., CNC [122], yields unsatisfactory OOD generalizationperformance, which further implies the necessity of the decomposition in CIGA.
Moreover, the architecture of CIGA can have multiple other implementations for both the featurizerand classifier, such as identifying G c at the latent space [85, 86].
Since we cannot enumerate everypossible implementation, in this work we choose interpretable GNN architectures as a prototypevalidation for CIGA and leave more sophisticated architectures as future works.
In particular, whenoptimized with ERM objective, CIGA can be reduced to interpretable GNNs .
However, merelyusing interpretable GNNs such as ASAP [77], GIB [119] or DIR [103] cannot yield satisfactory OODperformance.
As shown in Table 1(b) and discussed in Appendix.
D.4, GIB can only work for FIIF,while DIR cannot yield OOD guarantees for neither FIIF and PIIF SCMs.
These results are alsoempirically validated in the experiments.
We provide more detailed discussions in Appendix B.
Empirical Studies
We conduct extensive experiments with 16 datasets to verify the effectiveness of CIGA.
Datasets.
We use the SPMotif datasets from DIR [103] where artificial structural shifts and graph sizeshifts are nested (SPMotif-Struc).
Besides, we construct a harder version mixed with attribute shifts(SPMotif-Mixed).
To examine CIGA in real-world scenarios with more complicated relationshipsand distribution shifts, we also use DrugOOD [40] from AI-aided Drug Discovery with Assay,
Scaffold, and Size splits, convert the ColoredMNIST from IRM [4] using the algorithm from Knyazevet al.
[46] to inject attribute shifts, and split Graph-SST [120] to inject degree biases.
To compare withprevious specialized OOD methods for graph size shifts [112, 11], we use the datasets in Bevilacquaet al.
[11] that are converted from TU benchmarks [66].
More details can be found in Appendix G.1.
Baselines and our methods.
Besides the ERM, we also compare with SOTA interpretable GNNs,
GIB [119], ASAP Pooling [77], and DIR [103], to validate the effectiveness of the optimizationobjective in CIGA.
We use the same selection ratio (i.e., s c ) for all models.
Moreover, to validatethe effectiveness of the decomposition in CIGA, we compare CIGA with SOTA OOD objectivesincluding IRM [4], v-Rex [49] and IB-IRM [2], for which we apply random environment partitionsfollowing [23].
We also compare CIGA with EIIL [23] and CNC [122] that do not require environment labels, where CNC [122] has a more sophisticated contrastive sampling strategy for combatingsubpopulation shifts.
More implementation and comparison details are deferred to Appendix G.2.
Evaluation.
We report the classification accuracy for all datasets, except for DrugOOD datasetswhere we use ROC-AUC following [40], and for TU datasets where we use Matthews correlationcoefficient following [11]. We repeat the evaluation multiple times, select models based on the
[Table]
validation performances, and report the mean and standard deviation of the corresponding metric. Foreach dataset, we also report the “Oracle” performances that run ERM on the randomly shuffled data.
OOD generalization performance on structure and mixed shifts.
In Table 1, we report thetest accuracy of each method, where we omit GIB due to its poor convergence.
Different biasesindicate different strengths of the distribution shifts.
Although the training accuracy of most methodsconverges to more than 99% , the test accuracy decreases dramatically as the bias increases and asmore distribution shifts are mixed, which concurs with our discussions in Sec.
2.3 and Appendix D.
Due to the simplicity of the task as well as the relatively high support overlap between training andtest distributions, interpretable GNNs and OOD objectives can improve certain OOD performance,while they can have high variance since they donot have OOD generalization guarantees.
In contrast,
CIGAv1 and CIGAv2 outperform all of the baselines by a significant margin up to 10% with lowervariance , which demonstrates the effectiveness and excellent OOD generalization ability of CIGA.
[Table]
OOD generalization perfor-
[Table]
mance on realistic shifts.
In
Table 2 and Table 3, we examine the effectiveness of CIGA inreal-world data and more complicated distribution shifts.
Bothaveraged accuracy and ranks arereported because of the datasetheterogeneity.
Since the tasks areharder than synthetic ones, interpretable GNNs and OOD objectives perform similar to or evenunder-perform the ERM baselines, which is also consistentto the observations in non-linearbenchmarks [34, 40].
However,both CIGAv1 and CIGAv2 consistently and significantly outperform previous methods, including previous specialized methods Γ
GNNs [11] for combating graph size shifts, demonstrating the generality and superiority of CIGA.
Comparisons with advanced ablation variants.
As discussed in Sec. 3.3, CIGA can be reduced tointerpretable GNNs and contrastive learning approaches.
However, across all experiments, we canobserve that neither the advanced interpretable GNNs (DIR) nor sophisticated contrastive objectiveswith specialized sampling strategy (CNC) can yield satisfactory OOD performance, which servesas strong evidence for the necessities of the decomposition as well as the objective in CIGA.
Furthermore, although CIGAv1 can outperform CIGAv2 when we may have a relatively accurates c , the improvements in CIGAv1 are not as stable as CIGAv2 or even unsatisfactory when theassumption is violated.
This phenomenon also reveals the advances of CIGAv2 in practical scenarios.
Conclusions
We studied the OOD generalization on graphs via graph classification, and propose a new solution
CIGA through the lens of causality.
By modeling potential distribution shifts on graphs with SCMs,we generalized and instantiated the invariance principle to graphs, which was shown to have promisingtheoretical and empirical OOD generalization ability under a variety of distribution shifts.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Sec.
B.4 in the appendix.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Sec.
Ain the appendix.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Sec.
C andSec.
E in the appendix.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Sec.
E in theappendix.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[No] A anonymous
URL will be provided during the discussion period.
Full code will be released after thepaper is accepted.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Sec.
G.(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Sec.
G.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Sec.
G.3.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes](b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A]
The data used are all publicly available datasets.
(e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
We didn’t conduct research with human subjects.
(b) Did you describe any potential participant risks, with links to Institutional ReviewBoard (IRB) approvals, if applicable?
[N/A]
We didn’t conduct research with human
subjects.
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A] We didn’t conduct research with humansubjects. A Broader
Impacts
{+B More Discussions on Related Works and Future Directions 
 B.1 More backgrounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 Detailed related work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .+}
B.3 More discussions on connections of CIGA with existing work . . . . . . . . . . . B.4 Discussions on limitations of CIGA and future directions . . C Full Structural Causal Models on Graph Generation [-
-]C.1 Discussions on specific cases of the SCMs . . . . . . . . . . . . . . . . . . . . . .
D More Details about Failure Case Studies in Sec. 2.
D.1 More empirical details about failure case study in Sec.
D.2 More discussions about failure case study in Sec.
D.3 More empirical results about failure case study in Sec.
D.4 Theoretical discussions for failure case study in Sec.
D.5 Challenges of OOD generalization on graphs. . . . . . . . . . . . . . . . . . . . .
E Theory and Discussions
[-and Discussions 
-]
E.1 More discussions on the Definition 2.5 for Invariant GNNs . .
E.2 Proof for theorem 3.1 (i) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.3 Proof for theorem 3.1 (ii) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Details of Prototypical CIGA Implementation
G Detailed Experimental Settings
G.1 Details about the datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.2 Training and Optimization in Experiments . . . . . . . . . . . . . . . . . . . . . .
G.3 Software and Hardware . . . .
G.4 Additional Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.5 Interpretation Visualization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A Broader Impacts
Considering the wide applications and high sensitivity of GNNs to distribution shifts and spuriouscorrelations, it is important to develop GNNs that are able to generalize to OOD data, especially forrealistic scenarios such as AI-aided Drug Discovery where OOD data are ubiquitous.
By formulating
OOD generalization problem on graphs using causality, our work can serve as an initiate step towardstackling OOD generalization problem on graphs, with the hope to empower GNNs for broader applications and social benefits.
Besides, this paper does not raise any ethical concerns.
This study does notinvolve any human subjects, practices to data set releases, potentially harmful insights, methodologiesand applications, potential conflicts of interest and sponsorship, discrimination/bias/fairness concerns,privacy and security issues, legal compliance, and research integrity issues.
B More Discussions on Related Works and Future DirectionsB.1 More backgrounds
We give more 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 featuredimension of d .
In graph classification, we are given a set of N graphs { G i } Ni =1 ⊆ G and their labels{ Y i } Ni =1 ⊆ Y =
R c from c classes.
Then, we train a GNN ρ ◦ h with an encoder h : G → R h thatlearns a meaningful representation h G for each graph G to help predict their labels y G
= ρ ( h G ) witha downstream classifier ρ : R h → Y .
The representation h G is typically obtained by performingpooling with a READOUT function on the learned node representations:
[Equation]
(6) where the READOUT is a permutation invariant function (e.g., SUM, MEAN) [106, 114, 69, 106,19 , 67], and h ( K ) u stands for the node representation of u ∈ V at K -th layer that is obtained byneighbor aggregation:
[Equation]
where N ( u ) is the set of neighbors of node u , σ ( · ) is an activation function, e.g., ReLU, and a ( · ) isan aggregation function over neighbors, e.g., MEAN.
Invariant Learning.
Invariant learning typically considers a supervised learning setting based on thedata D
= {D e } e collected from multiple environments E all , where D e = { G ei , y ei } is the dataset fromenvironment e ∈ E all .
( G ei , y ei ) from a single environment e are considered as drawn independentlyfrom an identical distribution P e .
The goal of OOD generalization is to train a GNN ρ ◦ h : G → Ywith data from training environments D tr = {D e } e ∈E tr ⊆E all , and generalize well to all (unseen)environments, i.e., to minimize:
[Equation]
where R e is the empirical risk under environment e [96, 75, 4].
More details can be referred in [2].
B.2 Detailed related work
GNN Explainability.
Works in GNN explainability aim to find a subgraph of the input graphas the explanation for the prediction of a GNN model [115, 120].
Although some may leveragecausality in explanation generation [53], they mostly focus on understanding the predictions of
GNNs in a post-hoc manner instead of OOD generalization.
Recently there are two works aiming toprovide robust explanations under distribution shifts, i.e., GIB [119] and DIR [103], and both of themfocus on tackling FIIF spurious correlations (Assumption C.2).
The theoretical guarantees of GIBfollows the theory of information bottleneck [94], while GIB can not solve PIIF spurious correlations(Assumption C.3).
As both FIIF and PIIF widely exist in realistic scenarios, failing to solve either ofthem could result in severe performance degradation in practice [4, 2, 5, 70].
While for DIR, thoughas a generalization of Chang et al.
[14] to graphs, can not provide any theoretical guarantees under
FIIF spurious correlations as shown in Appendix D.4, nor under PIIF spurious correlations.
GNN Extrapolation.
Recently there is a surge of attention in improving the extrapolation ability of GNNs and apply them to various applications, such as mathematical reasoning [83, 84],physics [6, 82], and graph algorithms [93, 97, 107, 104].
Xu et al.
[109] study the neural networkextrapolation ability from a geometrical perspective.
Han et al.
[36] improve OOD drug discoveryby mitigating the overconfident misprediction issue.
Knyazev et al.
[46], Yehudai et al.
[112] focuson the extrapolation of GNNs in terms of graph sizes, while making additional assumptions onthe knowledge about ground truth attentions and access to test inputs.
Bevilacqua et al.
[11] studythe graph size extrapolation problem of GNNs through a causal lens, while the induced invarianceprinciple is built upon assumptions on the specific family of graphs.
Different from these works, weconsider the GNN extrapolation as a causal problem, establish generic SCMs that are compatiblewith several graph generation models, as well as, more importantly, different types of distributionshifts.
Hence, the induced the invariance principle and provable algorithms built upon the SCMs inour work can generalize to multiple graph families and distribution shifts.
Additionally, Wu et al.
[102] propose causal models as well as specialized objectives to extrapolatenodes with different neighbors.
However, their formulation is limited to node classification task andspecific spurious correlation type.
In contrast, the induced invariance principle in Wu et al.
[102],can be seen as a extension of CIGA for node classification, where we cab identify an invariantsubgraph from the K -hop neighbor graph of each node, and making predictions based on it, i.e.,
Y ⊥⊥ E | G ego c ⊆ G ego u for node u .
We leave specific formulation and implementation to future works.
Causality and OOD Generalization.
Causality comes to the stage for demystifying and improvingthe huge success of machine learning algorithms to further advances [74, 85, 86].
One of the mostwidely applied concept from causality is the Independent Causal Mechanism (ICM) that assumesconditional distribution of each variable given its causes (i.e., its mechanism) does not inform orinfluence the other conditional distributions [73, 76].
The invariance principle is also induced fromthe ICM assumption.
Once proper assumptions about the underlying data generation process via
Structural Causal Models (SCM) are established, it is promising to apply the invariance principleto machine learning models for finding an invariant representation about the causal relationshipbetween the underlying causes and the label [75, 4].
Consequently, models built upon the invariantrepresentation can generalize to unseen environments or domains with guaranteed performance [75,78 , 4, 80, 10, 48, 34, 49, 23, 2].
The arguably first formulation of invariance principle was introducedby Peters et al.
Arjovsky et al.
[4] propose a novel formulation of learning causal invariance inrepresentation learning, i.e., IRM, show how it connects with existing areas such as distributionalrobust optimization [71] and generalization [121], and prove its effectiveness in addressing PIIFspurious correlations (Assumption C.3).
However, in practice, both PIIF and FIIF (Assumption C.2)can appear in data, while IRM can fail in these cases [5, 70].
Ahuja et al.
[2] then propose to addinformation bottleneck criteria into the IRM formulation to address the issue.
However, their resultsare restricted to linear regime and also require environment partitions to distinguish the sources ofdistribution shifts.
Recently, Creager et al.
[23] and Lin et al.
[54] propose new OOD objectives torelieve the needs for environment partitions, but limited to PIIF spurious types and linear regime.
In parallel invariant learning approaches, Sagawa* et al.
[80] propose to regularize the worst groupin group distributionally robust optimization (GroupDro).
Zhang et al.
[122] propose a contrastiveapproach to tackle GroupDro when the group partitions are not available.
However, minimizing thegap between worst group risk and averaged risk can not yield a OOD generalizable predictors in ourcircumstances.
Besides, traditional approaches to tackle OOD generalization also include Domain
Adaption, Transfer Learning and Domain Generalization[78, 21, 31, 92, 52, 27, 60, 99], which aimto learn the class conditional invariant representation shared across source domain and target domain.
However, they all require a stronger assumption on the availability of target domain data or theground truth predictors [34, 2], hence are not able to yield predictors with OOD generalizationguarantees.
We refer interested readers to Pearl [74], Sch¨olkopf [85], Sch¨olkopf et al.
[86] for anin-depth understanding, and Gulrajani and Lopez-Paz [34], Ahuja et al.
[2] for a thorough overview.
B.3 More discussions on connections of CIGA with existing work
Although primarily serving for graph OOD generalization problem, our theory complements theidentifiability study on graphs through contrastive learning, and aligns with the discoveries in theimage domain that contrastive learning learns to isolate the content ( C ) and style ( S ) [125, 50].
Moreover, our results also partially explain the success of graph contrastive learning [117, 59, 118],where GNNs may implicitly learn to identify the underlying invariant subgraphs for prediction.
On expressivity of graph encoder in CIGA.
The expressivity of CIGA is essentially constrainedby the encoders embedded for learning graph representations.
During isolating G c from G , if theencoder can not differentiate two isomorphic graphs G c and G c ∪ G ps where G ps ⊆
G s , then thefeaturizer will fail to identify the underlying invariant subgraph.
Moreover, the classifier will also failif the encoder can not differentiate two non-isomorphic G c s from different classes.
Thus, adoptingmore powerful graph representation encoders into CIGA can improve the OOD generalization.
On CIGA and graph information bottleneck.
Under the FIIF assumption on latent interaction, theindependence condition derived from causal model can also be rewritten as Y ⊥⊥ S | C (similar to thatin DIR [103] as they also focus on FIIF), which further implies Y ⊥⊥ S | (cid:98) G c .
Hence it is natural to use
Information Bottleneck (IB) objective [94] to solve for G c :
[Equation]
which explains the success of many existing works in finding predictive subgraph through IB [119].
However, the estimation of I ( (cid:98) G c , G ) is notoriously difficult due to the complexity of graph, whichcan lead to unstable convergence as observed in our experiments.
In contrast, optimization withcontrastive objective in CIGA as Eq.
5 induces more stable convergence.
On CIGA for node classifications.
As the task of node classification can be viewed as graphclassification based on the ego-graphs of a node, our analysis and discoveries can generalize to nodeclassification.
More specifically, the invariance principle for node classification can be implementedby identifying an invariant subgraph from the K -hop neighbor graph of each node, and makingpredictions based on it, i.e., Y ⊥⊥ E | G ego c ⊆
G ego u for node u [102].
B.4 Discussions on limitations of CIGA and future directions
Better graph generation modeling.
Compared to Bevilacqua et al.
[11], we do not specify a specificgraph family in the SCM for graph generation process.
Since our focus is to describe the potentialdistribution shifts with SCMs, in Assumption 2.1, we aim to build a SCM that is compatible to manygraph generation processes [88, 56, 116, 58].
However, it is often the case that practitioners havecertain inductive knowledge about the graph generation process, which may imply useful leads andinvariance in modeling the generation process [110, 30, 55].
In Appendix C.1, we provide an exampleabout incorporating the graphon [56] knowledge into the SCMs, which derives similar solutions as inthe literature [112, 11].
Therefore, we believe it is promising to leverage more additional knowledgefor more precise graph generation modeling and better OOD generalization on graphs.
Better contrastive sampling.
Typical contrastive or graph contrastive learning approaches leverageaugmentation techniques as well as sophisticated sampling strategies during the positive or negativepairs selection [20, 81, 95, 117, 118].
A better augmentation or sampling strategy can benefit the
OOD generalization in general as shown by K¨ugelgen et al.
[50] and Zhang et al.
Sinceour implementation of CIGA in this work aims to verify the theoretical findings, we do not applysophisticated augmentation or sampling during the sampling while simply using the supervisedcontrastive approach [42].
Thus, it is promising to leverage better augmentation and contrastivestrategy to improve the generalization ability in CIGA.
More sophisticated architectures/parameter tunning.
The CIGA framework introduced in Sec.
3can have multiple implementations.
We choose interpretable architectures in our experiments for thepurpose of concept verification.
Essentially, different architectures can have different advantages andlimitations.
For the interpretable GNNs used in our experiments, it can provide interpretability for theresults (as shown in Appendix G.5), but still requires more training time (as shown in Appendix G.4).
Therefore, it may not be applicable to some resource-limited scenarios such as Edge-AI.
Besides,the approximation may also be limited to the chosen architectures.
More sophisticated architecturescan be incorporated, such as identifying and disentangling G c at the latent space [85, 86].
Moreover,as shown in Appendix G.4, CIGA still requires certain additional tunning efforts for the objectives.
Hence we believe it is also a promising future direction to reduce the parameter tunning by leveragingbetter optimization techiniques [87, 18]
C Full Structural Causal Models on Graph Generation
Due to the space constraints in the main paper, we make some simplifications when giving the SCMson the graph generation process.
Hence in this section, supplementary to the graph generation processin Sec. 2.2, we provide full SCMs on the graph generation process in this section as shown in Fig. 4.
Formal descriptions are given as Assumptions C.1, C.2, C.3, C.4.
To begin with, we take a latent-variable model perspective on the graph generation process andassume that the graph is generated through a mapping f gen : Z → G , where Z ⊆ R n is the latentspace and G = ∪ ∞ N =1 { 0 , 1 } N × R N × d is the graph space.
Let E denote environments.
Followingprevious works [50, 2], we partition the latent variable from Z into an invariant part C ∈ C =
R n cand a varying part S ∈ S =
R n s , s.t., n = n c + n s , according to whether they are affected by E .
Similarly in images, C and S can represent content and style while E can refer to the locationswhere the images are taken [7, 123, 50].
While in graphs, C can be the latent variable that controlsthe generation of functional groups in a molecule, which can not be affected by the changes ofenvironments, such as species (or scaffolds), experimental environment for examining the chemicalproperty (or assays) [40].
On the contrary, the other latent variable S inherits environment-specificinformation thus can further affect the finally generated graphs.
Besides, C and S can have multipletypes of interactions at the latent space with environments E and labels Y , which will generatedifferent types of spurious correlations [2].
Assumption C.1 (Graph generation SCM) .
[Equation]
Specifically, the graph generation process is shown as Fig. 4(a).
The generation mapping f gen isdecomposed into f ( A,X ) c
gen , f G c gen , f ( A,X ) s gen , f G s gen and f G gen to control the generation of ( Z cA , Z cX ) , G c ,( Z sA , Z sX ) , G s , and G , respectively.
Given the variable partitions C and S at the latent space Z , theycontrol the generation of the adjacency matrix and features for the invariant subgraph G c and spurioussubgraph G s through two pairs of latent variables ( Z cA , Z cX ) and ( Z sA , Z sX ) , respectively.
Z cA and
Z sA will control the structure-level properties in the generated graphs, such as degrees, sizes, andsubgraph densities.
While Z cX and Z sX mainly control the attribute-level properties in the generatedgraphs, such as homophily.
Then, G c and G s are entangled into the observed graph G through f G gen .
It can be a simply JOIN of a G c with one or multiple G s , or more complex generation processescontrolled by the latent variables [88, 56, 116, 58, 11].
Note that since our focus is to describe thepotential distribution shifts with SCMs, in Assumption 2.1, we aim to build a SCM that is compatibleto many graph generation processes [88, 56, 116, 58]. In fact, in Appendix C.1, we showcase how our SCMs can generalize to specific graph families studied in the literature [11, 103, 102], when givenmore additional knowledge about the graph generation process.
Nevertheless, we believe integratingspecific graph generation processes and their implications to improving OOD generalization ongraphs would be a promising future direction, as discussed in Appendix B.4.
Due to the correlation between E and G , graphs collected from different environments can havedifferent structure-level properties such as degrees, graph sizes, and subgraph densities, as well asfeature-level properties such as homophily [46, 112, 11, 17].
Meanwhile, all of them can spuriouslycorrelated with the labels depending on how the underlying latent variables are interacted with eachothers.
The interaction types can be further divided into two axiom types FIIF and PIIF, as wellas the mixed one MIIF.
Previous OOD methods such as GIB [119] and DIR [103] mainly focuson FIIF case, while others such as IRM [4] mainly focuses on the PIIF case.
Evidences show thatfailing to model either of them when developing the OOD objectives can have serious performancedegenerations in practice [5, 70].
That is why we aim to model both of them in our solution.
Assumption C.2 (FIIF SCM) .
[Equation]
Assumption C.3 (PIIF SCM) .
[Equation]
S C
Z cX
Z cA
Z sX
Z sA
G c
G s
[Figure]
E
S Y
G
C (
[Figure]
) FIIF SCM
E
S Y
G
C (
[Figure]
) PIIF SCM
E
S
C
S 2 Y
G ( [Figure] ) MIIF SCM Assumption C.4 (MIIF SCM) .
[Equation]
As for the interactions between C and S at the latent space, we categorize the interaction modes into
Fully Informative Invariant Features (FIIF, Fig. 4(b)), and Partially Informative Invariant Features(PIIF, Fig. 4(c)), depending on whether the latent invariant part C is fully informative about label
Y , i.e., ( S, E ) ⊥⊥ Y | C .
It is also possible that FIIF and PIIF are entangled into a Mixed Informative
Invariant Features (MIIF,Fig. 4(d)).
We follow Arjovsky et al.
[4], Ahuja et al.
[2] to formulate the
SCMs for FIIF and PIIF, where we omit noises for simplicity [73, 76].
Since MIIF is built upon FIIFand PIIF, we will focus on the axiom interaction modes (FIIF and PIIF) in this paper, while most ofour discussions can be extended to MIIF or more complex interactions built upon FIIF and PIIF.
Among all of the interaction modes, f gen corresponds to the graph generation process in Assumption C.1.
f spu is the mechanism describing how S is affected by C and E at the latent space.
In
FIIF, S is directly controlled by C while in PIIF, indirectly controlled by C through Y , which canexhibit different behaviors in practice [2, 70].
Additionally, in MIIF, S is further partitioned into S 1and S 2 depending on whether it is directly or indirectly controlled by C , respectively.
Moreover,f inv : C → Y indicates the labeling process, which assigns labels Y for the corresponding G merelybased on C .
Consequently, C is better clustered than S when given Y [13, 15, 85, 86], which alsoserves as the necessary separation assumption for a classification task [68, 16, 64].
Assumption C.5 (Latent Separability) .
H ( C | Y ) ≤ H ( S | Y ) .
C.1 Discussions on specific cases of the SCMs
Although our primary focus in this work is to characterize general graph distribution shifts that couldhappen in practice without any additional knowledge about the underlying graph family, and derivethe corresponding solutions, our SCMs (Fig. 4) can generalize to specific cases studied in previousworks, when incorporating more inductive biases about the underlying graph family [11, 103, 102]. Specifically, we illustrate the specialized SCMs in Fig. 5 for the SCM studied in [11] which assumesthe graphs are generated following a graphon model [56].
When with the additional knowledge about the underlying graph generative model, the graph generation SCM (Fig. 4(a)) and the FIIF SCM (Fig. 4(b)) together generalizes to the graphon SCM studiedin [11].
We now give a brief description in the below.
Specifically, shown as in Fig.
5(a), C now is instantiated as a graphon model C W ∼ P ( C W ) , where
C W : [0 , 1] 2 → [0 , 1] is a random symmetric measurable function sampled from the set of allsymmetric measurable functions [56].
Besides, the label Y is determined according to C W .
Then,
C W will further control the generation of the adjcency matrix G c
=
A c through graphon generativeprocess:
[Equation]
where Z u,v is an independent uniform noises on [0 , 1] for each possible edge ( u, v ) in the graph.
Bascially
Z graphon SCM
Fig. 5(c).
[Figure]
from [11]
[Figure]
On the other hand, as S does not imply any information about Y in this case, it resembles the FIIF
SCM (Fig. 4(b)).
In other words, ( S, E ) ⊥⊥ Y | C still holds.
Moreover, the node attributes G s =
X sare generated jointly influenced by the environment E and the graphon C W through S :
[Equation]
which resembles the attribute generation in Fig. 5(c).
Then, both G c and G s are concatenated together.
In a simplistic case intuitively, we can regard G conly contains the edges in G and G s only contains the node attributes.
Since the graphon modelmainly controls the edge connection, the edge connection patterns, e.g., motif appearance frequencyor subgraph densities, acts as a informative indicator for the label Y .
In contrast, the node attributesand its numbers would be affected by the environments. A GNN model is prone to the changesof the environments if it overfits to some spurious patterns about the graph sizes or the attributes. While if the GNN model can leverage the connection patterns to make predictions, it remain invariantto the changes of environments, or the spurious patterns such as graph sizes and node attributes,which resembles the solutions derived in [112, 11]. Besides, it also partially explains why CIGA cangeneralize to OOD graphs studied in these works [112, 11].
[Figure] In addition to the graphon SCM, essentially, the SCM studied in [103] resembles the FIIF SCM, andthat of [102] resembles PIIF SCM, which also serves as partial evidence for the superiority OOD
[Figure]
generalization performances of CIGA. D More Details about Failure Case Studies in Sec. 2.3 In this section, we provide details on failure case studies in Sec. We first elaborate the empiricalevaluation setting where we construct a synthetic graph datasets to probe the behaviors of existingmethods in OOD generalization on graphs.
D.1 More empirical details about failure case study in Sec. 2.3
To begin with, we construct 3-class synthetic datasets based on BAMotif [57] and follow Wu et al.[103] to inject spurious correlations between motif graph and base graph during the generation.
Inthis graph classification task, the model needs to tell which motif the graph contains, e.g., “House”or “Cycle” motif, as shown in Fig.
We inject the distribution shifts in the training data whilekeeping the test data and validation data without the biases.
For structure-level shifts, we introducethe artificial bias based on FIIF, where the motif and the base graph are spuriously correlated with aprobability of various bias.
For mixed shifts, we additionally introduced attribute-level shifts basedon FIIF, where all of the node features are spuriously correlated with a probability of various bias.
The number of training graphs is 600 for each class and the number of graphs in validation and testset is 200 for each class.
More construction details are given in Appendix G.
For the GNN encoders, by default, we use 3 -layer GCN [45] with mean readout, a hidden dimensionof 64 , and JK jump connections [105] at the last layer.
During training, we use a batch size of 32 ,learning rate of 1 e − 3 with Adam optimizer [43], and batch normalization between hidden layers [39].
Meanwhile, to stabilize the training, we also use dropout [90] of 0 .
1 and early stop the training whenthe validation accuracy does not increase till 5 epoch after first 20 epochs.
All of the experiments arerepeated 5 times, and the mean accuracy as well as variance are reported and plotted.
When using IRM objective [4], as the environment partitions are not available, we generate 2 environments withrandom partitions.
D.2 More discussions about failure case study in Sec. 2.3
[Figure] In Fig. 7, 8, 9, 10, we investigate whether existing training objectives (ERM and IRM), addingmore message passing, as well as using expressive GNNs, can improve the OOD generalization ability on graphs.
Here we also provide aadditional discussion in complementary to thediscussions on OOD generalization performanceof ERM and IRM objectives in Sec. 2.
Can better architectures improve OOD generalization of GNNs?
Adding more message passing turns.
It is acommon practice in GNNs to denoise the signalsby aggregating more neighbors with higher layers, or enhance the expressive power with morepowerful readout functions [105, 106, 111].
Aggregating neighbor information with more layersto denoise the input signal, or enhancing the expressivity with more powerful readout functions,are two common choices in GNNs to improvethe generalization ability [105, 51, 106, 111].
However, in the experiments next, we empirically found that GCNs with more layers andmore powerful readout operations are still sensitive to distribution shifts.
In particular, stackingmore layers helps denoising certain shifts, whilethe OOD performance would drop more sharply when the bias increases.
Intuitively, if the spuriousfeatures from nodes cannot be eliminated by the denoising property of a deeper GNN, they wouldspread among the whole graph more widely, which in turn leads to stronger spurious correlations.
Besides, the spurious correlations would be more difficult to be disentangled if there are distributionshifts at both structure-level and attribute-level.
Since the node representations from hidden layerscan also encode graph topology features [106], distribution shifts introduced through Z s A and Z s Xwill doubly mix at the learned features.
In the worst case, the information about Z cA and Z cX couldbe partially covered by or even replaced by Z sA and Z sX .
This will make OOD generalization ofmessage passing GNNs trained through ERM much more difficult or even impossible.
Besides, asthe node representations of 1 ≤ i ≤ k -th layer can also encode graph topology features [106], which,if spuriously correlated with labels through Z sA and entangled with part of invariant node features,i.e., Z cX , in the worst case, can greatly improve the difficulty or even make the OOD generalizationimpossible for neighbor aggregation GNNs trained with ERM.
Using more expressive GNNs.
Previous results on the expressivity of GNNs show that GNNs arelimited to distinguish isomorphic graphs at most as 1-WL/2-WL test can distinguish [106].
Afterthat, many follow-up variants are proposed to improve the expressivity of GNNs [67].
However,if the labels are spuriously correlated with certain subgraphs, even the GNN has high expressivitycan still be prone to distribution shifts.
In a idealistic case, when classifying a graph with a highlyexpressive GNN, it reduces to the linear or discrete feature case on the Euclidean regime.
In this case,there exists many evidences showing that neural networks can fail to generalize to OOD data withouta proper objective [7, 24, 4, 80, 10, 49, 23, 48, 2].
Empirically, we use k -GNNs [65] to verify theintuition and observe similar failures for this provably more expressive GNN as basic GNN variants.
[Figure]
[Figure] [Figure] [Figure]
[Figure]
To explore the behaviors of aforementioned methods against complicated distribution shifts on graphs,we first modify construction method in Wu et al.
[103] to construct dataset for Fig.
7, where only
FIIF structure-level spurious correlations are injected.
Then we also inject FIIF attribute-level shifts,by setting the node attributes to constant vectors which is spuriously correlated with the labels.
Furthermore, in Fig. 8, graph size shifts are added, which is exactly the SPMotif datasets used in
DIR [103].
Besides, in Fig.
9, we can also change the FIIF attribute-level shifts to PIIF attribute-levelshifts, where we flip the labels by a probability of 5% and let the flipped label to be spuriouslycorrelated with the node features, following the PIIF SCM in Fig. 4.
Graph size shifts can also beinjected in this case, shown as Fig.
Next, we summarize our findings from the experiments.
Observation I: All existing methods are sensitive to distribution shifts.
From the Fig.
7, 8, 9, 10,we can observe that all GNNs are sensitive to distribution shifts.
As the intensity of spuriouscorrelation grows, GNNs are more likely to overfit to shortcuts presented either in the structure-levelor attribute-level, which is similar to general deep learning models [33].
Observation II: Higher variance also indicates unstable OOD performance.
Although GNNsshow certain robustness against single distribution shifts, e.g., performances do not decrease sharplyat the beginning in Fig.
7, when the spurious correlation grows stronger, the OOD performancebecome more unstable , e.g., higher variance.
The reason is that, GNNs sometimes can directlylearn about the desired information at some random initializations, since the task is relatively simplecompared to reality.
Hence the performance will be highly sensitive to the quality of initialized pointsat the beginning.
Consequently, the performances from multiple runs would exhibit high variance.
However, when the task becomes more difficult, GNNs will consistently be prone to distributionshifts, and the variance will be smaller, as shown in experiments (Sec. 4).
Observation III: Entangling more distribution shifts can degenerate more GNN performance.
As implied by the graph generation SCMs in Fig.
4, distribution shifts can happen at both structurelevel and attribute-level, and each of them can have different type of spurious correlation with the label.
In Fig.
7, we can find that, when the attribute-level distribution shifts are mixed, the performancewill be worse and more unstable.
When the graph size shifts are mixed, this phenomenon will bemore obvious, as shown in Fig.
This phenomenon also verifies the observations in Knyazev et al.[46] that attention mechanism in GNN is also sensitive to graph size shifts and can hardly learn thedesired attention distributions without further guidance.
Moreover, when the structure-level andattribute-level shifts have different spurious correlation types, i.e., when FIIF structure-level shiftsand PIIF attribute-level shifts are both presented, the performance drop will be more serious, bycomparing Fig.
7 to Fig.
9, as well as Fig.
8 to Fig.
Observation IV: Using more powerful architectures can not improve the OOD performance.
From the sub-figures (b) and (c) in Fig.
7, 8, 9, 10, we can also observe that neither adding moremessage passing turns nor using more expressive GNN architectures can be immune to distributionshifts.
On the contrary, they also exhibit similar behaviors like basic GNN architectures.
Specifically,adding more message passing runs show certain robustness against distribution shifts since theyare more likely to learn the desired information during the optimization [108].
However, when theintensity of spurious correlation grows stronger, deeper GNNs are more likely to overfit to shortcutshence their performances will drop more sharply.
On the other hand, using provably more expressive
GNN architectures can not improve the OOD performance, either.
In Fig.
7, 8, 9, 10 we use 1 - 2 - 3GNN following the algorithm of k -GNNs which is provably more expressive than 2 -WL test [65].
When there are no graph size shifts, k -GNNs will have higher performance at the beginning.
Whenthere are graph size shifts, k -GNNs will have a lower initial performance at the beginning.
Then, asthe spurious strength grows, k -GNNs can suddenly become seriously unstable, though k -GNNs canhave higher averaged performance, which reflects unsatisfactory OOD performance as Observation IIimplies.
When the intensity of spurious correlations grows even stronger, similar to deeper GNNs,
OOD performances of k -GNNs will be more unstable and go down to similar level as that of normal
GNN architectures.
Hence, it calls for better optimization objectives as well as a suitable architecturesto help improve the OOD generalization performance.
Beyond the empirical studies in previous section, we aim to accompany more formal discussions forexplaining the failures of existing optimization objectives and architectures in the next sections.
D.4 Theoretical discussions for failure case study in Sec. 2.
A motivating example.
To begin with, we follow Ahuja et al.
[2] to introduce a formal example onthe failures of GNNs optimized with ERM or IRM [96, 4] via a linear binary classification problem:
Definition D.1 (Linear classification structural equation model (FIIF)) .
[Equation]
where w ∗ inv ∈ R n c with ∥ w ∗ inv ∥
= 1 is the labeling hyperplane, C ∈ R n c , S ∈ R n s are thecorresponding invariant and varying latent variables, N is Bernoulli binary noise with a parameter ofq and identical across all environments, ⊕ is the XOR operator, S is invertible.
Given data generation process as Assumption C.1, and latent space interaction as Assumption C.2 or
C.3, and strictly separable invariant features 2.4, consider a k -layer linearized GNN ρ ◦ h using meanas READOUT for binary graph classification, if ∪ e ∈E te supp ( P e ) ̸⊆ ∪ e ∈E tr supp ( P e ) :
(i) For graphs features generated as Definition D.1, ρ ◦ h optimized with ERM or IRM will fail
to generalize OOD (Eq. 8) almost surely;
(ii) For graphs with more than two nodes, globally same node features generated as Defini-
tion D.1, and graph labels that are the same as global node labels, ρ ◦ h optimized with ERM
or IRM will fail to generalize OOD (Eq.
8) almost surely;
For graph classification, if the number of nodes is fixed to one, it covers the linear classification asabove.
When ∪ e ∈E te supp ( P e ) ̸⊆ ∪ e ∈E tr supp ( P e ) , it implies the S from training environments E tr doesnot cover S from testing environments, while C can be covered.
Moreover, the condition of strictlyseparable training data now can be formulated as min C ∈∪ e ∈E tr ( C [Equation] G e ) sgn ( w ∗ inv · C )( w ∗ inv · C ) > 0 .
Recall that ERM trains the model by minimizing the empirical risk (e.g., 0-1 loss) over all trainingdata, and IRM formulates OOD generalization as:
[Equation]
However, both ERM and IRM can not enable OOD generalization, i.e., finding the ground truth w ∗ inv ,following the Theorem 3 from Ahuja et al.
Theorem D.2 (Insufficiency of ERM and IRM) .
Suppose each e ∈ E all follows Definition.
D.1, C arestrictly separable, bounded and satisfy the support overlap between E tr and E te , and S are bounded,if S does not support the overlap, then both ERM and IRM fail at solving the OOD generalizationproblem.
The reason is that, when C from all environments are strictly separable, there can be infinite many
Bayes optimal solutions given training data { G e , y e } e ∈E tr , while there is only one optimal solutionthat does not rely on S .
Hence, the probability of generalization to OOD (finding the optimal solution)tends to be 0 in probability.
As for case (ii), when the GNN uses mean readout to classify more than one node graphs, assumingthe graph label is determined by the node label and all of the nodes have the same label that aredetermined as Definition D.1, then GNN optimized with ERM and IRM will also fail because of thesame reasons as case (i).
Discussions on the failures of previous OOD related solutions.
First of all, for IRM or similarobjectives [80, 49, 2, 9] that require environment information or non-trivial data partitions, they canhardly be applied to graphs due to the lack of such information.
The reason is that obtaining suchinformation can be expensive due to the abstraction of graphs.
Moreover, as proved in Theorem 5.1of Rosenfeld et al.
[79], when there is not sufficient support overlap between training environmentsand testing environments, the IRM or similar objectives can fail catastrophically when being appliedto non-linear regime.
The only OOD objective EIIL [23] that does not require environment labels,also rely on similar assumptions on the support overlap.
We also empirically verify their failingbehaviors in our experiments.
Moreover, since part of explainability works also try to find a subset of the inputs for interpretableprediction robustly against distribution shifts.
Here we also provide a discussion for these works.
Thefirst work following this line is I NV RAT [14], which develops an information-theoretic objective (were-formulate it to suit with OOD generalization problem on graphs):
[Equation]
However, it also requires extra environment labels for optimization that are often unavailable ingraphs.
Besides, the corresponding assumption on the data generation for guaranteed performance isessentially PIIF if applied to our case, while it can not provide any theoretical guarantee on FIIF.
We also notice a recent work, DIR [103], as a generalization of I NV RAT to graphs while studying
FIIF spurious correlations, that proposes an alternative objective which does not require environmentlabel:
[Equation]
However, the theoretical justification established for DIR (Theorem 1 to Corollary 1 in Wu et al.
[103])essentially depends on the quality of the generator g which can be prone to spurious correlations.
Thus, DIR can hardly provide any theoretical guarantees when applied to our case, neither for FIIFnor PIIF.
In experiments, we empirically find the unstable and relatively high sensitivity of DIR tospurious correlations, which verifies our finding.
More details about empirical behaviors of DIR canbe found in Appendix G.
In contrast to DIR, GIB [119] that focuses on discovering a informative subgraph for explanation,essentially can provide theoretical guarantees for FIIF spurious correlations.
Theoretically, (wecopy the discussion in Appendix F here to provide an overview of relationships between GIB and
DIR.)
Under the FIIF assumption on latent interaction, the independence condition derived fromcausal model can also be rewritten as Y ⊥⊥ S | C (similar to that in DIR [103] as they also focuson FIIF), which further implies Y ⊥⊥ S | (cid:98) G c .
Hence it is natural to use Information Bottleneck (IB)objective [94] to solve for G c :
[Equation]
(13) which explains the success of many existing works in finding predictive subgraph through IB [119].
However, the estimation of I ( (cid:98) G c , G ) is notoriously difficult due to the complexity of graph, whichcan lead to unstable convergence as observed in our experiments.
In contrast, optimization withcontrastive objective in CIGA as Eq.
5 induces more stable convergence.
D.5 Challenges of OOD generalization on graphs.
From the aforementioned analysis, we can summarize some key challenges revealed by the failuresof both existing optimization objectives and GNN architectures.
In particular, we are facing two mainchallenges a) Distribution shifts on graphs are more complicated where different types of spuriouscorrelations can be entangled via different graph properties; b) Environment labels are usually notavailable due to the abstract graph data structure.
E Theory and Discussions
In this section, we provide proofs for propositions and theorems mentioned in the main paper.
E.1 More discussions on the Definition 2.5 for Invariant GNNs
Definition 2.5 is motivated by applying the invariance principle to the established SCMs in Sec. 2.2,following the literature of invariant learning [75].
In this section, we will present Proposition E.2 and Proposition E.3 to illustrate how satisfying the minmax objective in Definition E.1 is equivalent toidentifying the underlying invariant subgraph G c that contains all of the information about causalfactor C in G , under both FIIF and PIIF SCMs (Fig. 2(b) and Fig. 2(c)). Definition E.1 (Invariant GNN) . Given a set of graph datasets {D e } e and environments E all thatfollow the same graph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutationinvariant graph encoder h : G → R h and a downstream classifier ρ : R h → Y , ρ ◦ h is an invariant
GNN if it minimizes the worst case risk among all environments, i.e., min
[Equation]
First, we show that using the invariant subgraphs G c to predict Y can satisfy the minmax objectivemin max e ∈E all R e in Proposition E.2.
Proposition E.2.
Let G c denote the subgraph space for G c , given a set of graphs with their labels
D =
{ G ( i ) , y ( i ) }
Ni =1 and E all that follow the graph generation process in Sec.
2.2 (or Sec.
C), a GNNρ ◦ h : G c → Y that takes G c of G as the input to predict Y , and solves the following objective cangeneralize to OOD graphs, i.e., solving the minmax objective in Def.
E.1:
[Equation]
where R G c is the empirical risk over { G ( i ) c , y ( i ) } Ni =1 and G ( i ) c is the underlying invariant subgraph
G c for G ( i ) .
Proof.
We establish the proof with independent causal mechanism (ICM) assumption in SCM [73, 76].
In particular, given the data generation assumption, i.e., for both FIIF (Assumption 2.2) and PIIF(Assumption 2.3), we have: ∀ e,
[Equation]
where we use ICM for the first three equalities.
From Eq.
14, it suffices to know P ( Y | G c ) isinvariant across different environments.
Hence, a GNN predictor ρ ◦ h : G c → Y optimizedwith empirical risk given G c , essentially minimizes the empirical risk across all environments, i.e.,
[Equation]
Thus, if ρ ◦ h solves min R G c , it also solves min max R e , hence it elicits ainvariant GNN predictor according to Definition.
E.1.
Besides, we show in Proposition E.3 that only using the underlying invariant subgraphs G c to makepredictions can satisfy the minmax objectives.
Or equivalently, a GNN predictor solving the minmaxobjective can only rely on the underlying invariant subgraph G c to predict Y .
Proposition E.3.
Given a set of graph datasets {D e } e and environments E all that follow the samegraph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutation invariantgraph encoder h : G → R h and a downstream classifier ρ : R h
→ Y , ρ ◦ h that minimizes theworst case risk among all environments, i.e., min max e ∈E all R e , can not rely on any part of G s , i.e.,ρ ◦ h ( G ) ⊥⊥ G s .
Proof.
The proof for Proposition E.3 is straightforward.
Assuming that ρ ◦ h ( G ) ̸⊥⊥ G s , as E isinfluenced by the changes of E through S in both FIIF and PIIF SCMs (Fig. 2(b) and Fig.
2(c)),then
ρ ◦ h ( G ) ̸⊥⊥ E as well.
Consequently, there exists some graph G corresponding to G c , G es andρ ◦
h
( G ) =
Y under an environment e , such that we can always find a proper e ′ to make ρ ◦ h ( G ) ̸
= Y .
In contrast, the prediction of a GNN that satisfies
ρ ◦ h ( G ) ⊥⊥ G s remains invariant against arbitrarychanges of environments. Thus, it leads to a contradiction to the condition that min max e ′ ∈E all R e ′ . Therefore, a GNN that solves min max e ∈E all R e must satisfy ρ ◦ h ( G ) ⊥⊥ G s . Combining Proposition E.2 and Proposition E.3, we are highly motivated to find the underlyinginvariant subgraphs to make predictions about the original graphs, which converges to Eq. 1.
Tackling
Eq.
1 under the unavailability of E brings us two variants of CIGA solutions, as illustrated in
Section 3. E.2 Proof for theorem 3.1 (i) Theorem E.4 (CIGAv1 Induces Invariant GNNs) . Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c genin Assumption 2.1 are invertible, (b) samples from each training environment are equally distributed,i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr , if ∀ G c , | G c | =
s c , then a GNN f c ◦ g solves Eq.
4, is an invariant GNN(Def. 2.5).
Proof.
We re-write the objective as follows:
[Equation]
where (cid:98) G c
= g ( G ) , (cid:101) G c = g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G and G have the same label.
The proof of Theorem E.4 is essentially to show the estimated (cid:98) G c through Eq.
15 is the underlying
G c , then the maximizer of I ( (cid:98) G c ; Y ) in Eq.
15 can produce most informative and stable predictionsabout Y based on G , hence is an invariant GNN according to Definition.
In the next, we are going to take a information-theoretic view of the first term
I ( (cid:98) G c ; Y ) and thesecond term I ( (cid:98) G c ; (cid:101) G c | Y ) to conclude the proof.
We begin by introducing the following lemma:
Lemma E.5.
Given the same conditions as Thm.
E.4, I ( (cid:98) G c ; Y ) is maximized if and only if
I ( (cid:98) G c ; Y | E = e ) is maximized, ∀ e ∈ E tr .
The proof for Lemma E.5 is straightforward, given the condition that samples from each trainingenvironment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr .
Obviously, (cid:98) G c =
G c is amaximizer of I ( (cid:98) G c ; Y )
=
I ( C ; Y ) =
H ( Y ) , since f c gen : C → G c is invertible and C causes
Y .
However, there might be some subset G ps ⊆
G s from the underlying G s that entail the sameinformation about label, i.e., I ( G pc ∪ G ps ; Y )
=
I ( G c ; Y ) where (cid:98) G c = G pc ∪ G ps and G pc = G c ∩ (cid:98) G c .
For FIIF (Assumption 4(b)), it can not happen otherwise, let G lc
= G c − G pc , then we have:
[Equation]
where the second last equality is due to C → Y and the invertibility of f c gen : C → G c in FIIF, i.e.,
H ( Y | C ) =
H ( Y | G c )
=
H ( Y | G pc , G lc ) = 0 .
However, in PIIF, it can not hold since conditioning on
G pc , G ps can not determine Y , since S ̸⊥⊥ Y | C .
In other words, G s ̸⊥⊥ Y | G c , which means G s canimply some information about Y that is equivalent to I ( G lc ; Y | G pc ) .
To avoid the presence of spuriously correlated G s in (cid:98) G c , we will use the second term to eliminate it:
[Equation]
[Equation]
where (cid:98) G c
= g ( G ) , (cid:101) G c = g ( (cid:101) G ) are two positive samples drawn from the same class (i.e., conditionon the same Y ).
Since the all of the training environments are equally distributed, maximizing
I ( (cid:98) G c ; (cid:101) G c | Y ) is essentially maximizes I ( (cid:98) G c ,
E = ˆ e ; (cid:101) G c , E = ˜
e | Y ) , ∀ ˆ e, ˜ e ∈ E tr .
Hence, we have:
[Equation]
We claim
Eq.
can eliminate any potential subsets in the estimated (cid:98) G c .
[Equation]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure] G be the or equiva the are estimated [Figure] respectively.
We can also define sim [Figure]
ilar counterparts for G s : (cid:98) G ps , (cid:101) G ps are the partial (cid:98) G s , (cid:101) G s s contained in the estimated (cid:98) G c , (cid:101) G cwhile (cid:98) G ls , (cid:101) G ls are the left subsets (cid:98) G s , (cid:101) G s , respectively.
Recall the constraint that | G c | =
s c , hence if (cid:98) G pc ⊆ (cid:98) G c , then a corresponding (cid:98) G lc = (cid:98) G ∗ c − (cid:98) G pc will bereplaced by (cid:98) G ps in (cid:98) G c .
In this case, we have:
[Equation]
where the second equality is due to E = ˆ e is determined so that H ( E = ˆ e | (cid:98) G c , Y ) =
Compared
Eq.
19 to that when (cid:98) G c
=
(cid:98) G ∗ c , we have the entropy change as:
[Equation] [Equation]
Let ϵ
= H ( (cid:98) G ps | (cid:98) G pc , E = ˆ e, Y ) .
In a idealistic setting, when the noise of the generation process S :=f spu ( Y, E ) in PIIF tends to be 0 , i.e., ϵ → 0 , S is determined conditioned on E, Y , hence G s and anysubsets of G s are all determined.
Then, it suffices to know that in Eq.
20, H ( (cid:98) G ps | (cid:98) G pc , E =
ˆ e, Y )
=
0while
H ( (cid:98) G lc | (cid:98) G pc , E = ˆ e, Y ) > [-
-]0 since (cid:98) G lc can not be determined when given (cid:98) G pc , E = ˆ e, Y .
Thus,when some subset from G s is included in (cid:98) G c , it will minimize H ( (cid:98) G c , E
= ˆ e | Y ) .
In the next, we willshow how ϵ =
H
(cid:98) G ps | (cid:98) G pc , E = ˆ e, Y ) can be cancelled thus leading to a smaller H ( (cid:98) G c , E
= ˆ e | Y ) ,by considering the second term H ( (cid:98) G c , E
= ˆ e | (cid:101) G c , E = ˜
e, Y ) .
As for H ( (cid:98) G c , E = ˆ e | (cid:101) G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possiblecases into two:
(i) One of (cid:98) G c and (cid:101) G c contains some subset of G s ,
i
e., (cid:98) G c contains some (cid:98) G ps ⊆ (cid:98) G s ;(ii) Both (cid:98) G c and (cid:101) G c contain some (cid:98)
G ps ⊆ (cid:98) G s and (cid:101) G ps ⊆ (cid:101) G s , respectively.
For (i), we have:
[Equation]
G [Equation] , E = [Equation]
Y ) [Equation] Thus, H ( (cid:101) G ps | (cid:98) G c , E = ˜
e, Y, (cid:101) G lc , E = ˆ e ) =
and Thus, we have:
[Equation]
E [Equation]
ˆ
[Equation]
= ˜
[Equation]
Combing ∆ H ( (cid:98) G c , E = ˆ e | Y ) , we have:
[Equation]
where the last equality is because of the independence of (cid:98) G ps between (cid:101) G c , E =
˜
e conditioned on
Y, E = ˆ e .
Since conditioning will lower the entropy for both discrete and continuous variables [22,113], we have:
[Equation]
which implies the existence of (cid:98) G ps in (cid:98) G c will lower down the second term in Eq.
15 for the case (i).
For (ii), we have:
[Equation]
Similar to (i), H ( (cid:98) G ps | (cid:101) G pc , (cid:101) G ps , E = ˜
e, Y, (cid:98) G pc , E = ˆ e ) can be cancelled out with H ( (cid:98) G ps | (cid:98) G pc , E = ˆ e, Y ) .
Then, we have:
[Equation]
[Equation]
Since additionally conditioning on (cid:98) G ps in H ( (cid:98) G lc , E = ˆ e | (cid:101) G pc , (cid:101) G ps , E = ˜
e, Y ) can not lead to newinformation about (cid:98) G lc , we have:
[Equation]
[Equation]
which follows that ∆ I ( (cid:98) G c , E = ˆ e ; (cid:101) G c , E =
˜ e | Y ) < 0 .
To summarize, the ground truth G c is the only maximizer of the objective (Eq. 15), hence solving forthe objective (Eq. 15) can elicit an invariant GNN.
E.3 Proof for theorem 3.1 (ii)
Theorem E.6 (CIGAv2 Induces Invariant GNNs) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c genin Assumption 2.1 are invertible, (b) samples from each training environment are equally distributed,i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr , a GNN f c ◦ g solves Eq.
4, is an invariant GNN (Def. 2.5).
Proof.
We re-write the objective as follows:
[Equation]
(28) where (cid:98) G c
= g ( G ) , (cid:101) G c = g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G and G have the same label.
Similar to the proof for Theorem E.4, to prove Theorem E.6 is essentially to show the estimated(cid:98) G c through Eq.
28 is the underlying G c , hence the minimizer of Eq.
28 elicits an invariant GNNpredictor.
In the next, we also begin with a lemma:
Lemma E.7.
Given data generation process as Theorem E.6, for both FIIF and PIIF, we have:
I ( C ; Y ) ≥ I ( S ; Y ) , hence I ( G c ; Y ) ≥ I ( G s ; Y ) .
Proof for Lemma E.7.
For both FIIF and PIIF, Assumption 2.4 implies that H ( C | Y ) ≤ H ( S | Y ) .
Itfollows that I ( C ; Y )
=
H ( Y ) − H
( C | Y )
≥ H ( Y ) − H ( S | Y ) =
I ( S ; Y ) .
Then, since f G c gen : C → G cis invertible, we have I ( G c ; Y ) =
I ( C ; Y ) ≥ I ( S ; Y ) ≥ I ( G s ; Y ) .
Given Lemma E.7, we know (cid:98) G c at least contains some subset of the underlying G c , otherwise theconstraint I ( (cid:98) G s ; Y ) ≤ I ( (cid:98) G c ; Y ) will be violated since G c ⊆ (cid:98) G s in this case.
Assuming there are some subset of G s contained in (cid:98) G c , without loss of generality, we can divide allof the possible cases about (cid:98) G c into two: (i) (cid:98) G c only contains a subset of
the
underlying
G c ; (ii) (cid:98) G c contains a subset of the underlying G c as well as part of the
underlying G s
[Equation]
subsets in G c and G s are denoted as (cid:98) G lc and (cid:98) G ls , respectively. Similar notations are also applicable for the estimated (cid:101) G c from (cid:101) G . Before the discussion, let us inherit the notationsof subsets of G c , G s from the proof for Theorem E.4:
Let (cid:98) G ∗ c and (cid:101) G ∗ c be the ground truthinvariant subgraph G c s of (cid:98) G and (cid:101) G , (cid:98) G lc = (cid:98) G ∗ c −(cid:98) G c and (cid:101) G lc
=
(cid:101) G ∗ c − (cid:101) G c be the l eft (un-estimated)subsets from corresponding ground truth G c s,and (cid:98) G pc = (cid:98) G ∗ c − (cid:98) G lc and (cid:101) G pc = (cid:101) G ∗ c − (cid:101)
G lc be thecomplement, or equivalently, the p artial (cid:98) G ∗ c , (cid:101) G ∗ cthat are estimated in (cid:98) G c , (cid:101) G c , respectively.
Similarly, (cid:98) G ps , (cid:101) G ps are the partial (cid:98) G s , (cid:101) G s s containedin the estimated (cid:98) G c , (cid:101) G c while (cid:98) G ls , (cid:101) G ls are theleft subsets (cid:98) G s , (cid:101) G s , respectively.
First of all, case (i) cannot hold because, whenmaximizing I ( (cid:98) G c ; (cid:101) G c | Y ) , if ∃ (cid:98) G lc = (cid:98) G ∗ c − (cid:98) G c ,as shown in the proof for Theorem E.4, including (cid:98) G lc into (cid:98) G c can always enlarge I ( (cid:98) G c ; (cid:101) G c | Y ) ,while not affecting the optimality of I ( (cid:98) G s ; Y )
+
I ( (cid:98) G c ; Y ) by re-distributing (cid:98) G lc from (cid:98) G s to (cid:98) G c .
Consequently, (cid:98) G ∗ c must be included in (cid:98) G c , i.e., (cid:98) G ∗ c ⊆ (cid:98) G c .
As for case (ii), recall that, by the condition of equally distributed training samples from each trainingenvironment, maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) is essentially maximizing I ( (cid:98) G c , E
= ˆ e ; (cid:101) G c , E = ˜
e | Y ) ,∀ ˆ e, ˜ e ∈ E tr , hence, we have:
[Equation]
We claim Eq.
29 can eliminate any potential subsets in the estimated (cid:98) G c .
Similarly, we have:
[Equation]
where the second equality is due to E
= ˆ e is determined.
Compared to the case that (cid:98) G c
= (cid:98) G ∗ c , wehave:
[Equation]
Then, as for H ( (cid:98) G c , E = ˆ e | (cid:101) G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possiblecases into two:
(a) (cid:98) G c contains some (cid:98) G ps ⊆ (cid:98) G s ;(b) Both (cid:98) G c and (cid:101) G c contain some (cid:98) G ps ⊆ (cid:98) G s and (cid:101) G ps ⊆ (cid:101) G s , respectively.
For (a), we have: H (
[Equation] [Equation]
ps ,
[Equation]
Similarly to the proof for Theorem E.4, when considering ∆ I ( (cid:98) G c ; (cid:101) G c | Y ) , the effects of
H ( (cid:98) G p s | (cid:101) G c , E = ˜
e, Y, (cid:98) G ∗ c , E = ˆ e ) is cancelled out by H ( (cid:98) G p s | (cid:98) G ∗ c , E = ˆ e, Y ) .
Hence, we have:
[Equation]
For (b), we have:
[Equation]
Similarly, H ( (cid:98) G ps | (cid:101) G ∗ c , (cid:101) G ps , E = ˜
e, Y, (cid:98) G ∗ c , E = ˆ e ) = 0 can also be cancelled out by H ( (cid:98) G ps | (cid:98) G ∗ c , E =ˆ e, Y ) .
Moreover, for H ( (cid:98) G ∗ c | (cid:101) G ∗ c , (cid:101) G ps , E = ˜
e, Y, E
= ˆ e ) , (cid:101) G ps can not bring no additional informationabout (cid:98) G ∗ c , when conditioning on (cid:101) G ∗ c , Y, E = ˜
e .
Hence, we also have:
[Equation]
To summarize, when maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) , including any (cid:98) G ps ⊆ (cid:98) G ∗ s can not bring additionalbenefit while affecting the optimality of I ( (cid:98) G s ; Y )
+ I ( (cid:98) G c ; Y ) .
More specifically, when consideringthe changes to I ( (cid:98) G s ; Y ) + I ( (cid:98) G c ; Y ) ,
∀ G ps ⊆
G s , we have
I ( G − (cid:98) G ∗ c − G ps ; Y ) ≤ I ( G
[Equation]
(cid:98) G ∗
[Equation]
) , ∀ G ps ⊆
G s , while I ( Y ; (cid:98) G ∗ c , G ps ) =
I ( Y ; (cid:98) G ∗ c )
+ I
( Y ; (cid:98) G ps | (cid:98) G ∗ c ) , ∀ e ∈ E tr .
Consequently,
[Equation]
Hence, only the underlying G c is the solution to Eq. 28, which implies that solving for the objective(Eq. 28) can elicit an invariant GNN.
F Details of Prototypical CIGA Implementation
In fact, the CIGA framework introduced in Sec.
3 can have multiple implementations.
We chooseinterpretable architectures in our experiments for the purpose of concept verification.
More sophisticated architectures can be incorporated.
Experimental results in Sec. 4 also demonstrates that, evenequipped with basic GNN architectures, CIGA already has the excellent OOD generalization ability,hence it is promising to incorporate more advanced architectures from the prosperous GNN literature.
We now introduce the details of the architectures used in our experiments.
Recall that CIGAdecomposes a GNN model for graph classification into two modules, i.e., a featurizer:
g : G → G c and a classifier f c : G c → Y .
Specifically, for the implementation of Featurizer, we choose one of thecommon practices GAE [44] for calculating the sampled weights for each edge.
More formally, thesoft mask is predicted through the following equation:
[Equation]
[Figure]
GNNs need to classify graphs based on the specific motif (“House” or “Cycle”).
The featurizer g will extract an (orange colored) subgraph (cid:98) G c from each input for the classifier f c to predict the label.
The training objective of g is implemented in a contrastive strategy where the distribution of (cid:98) G c at the latent sphere will be optimized to maximize the intra-class mutual information.
With the identified invariant subgraph G c , the predictions made by classifier f c based on G c are invariant to distribution shifts;
If a sampling ratio s c is predetermined, we sample s c of total edges with the largest predictedweights as a soft estimation of (cid:98) G c .
Then, theestimated (cid:98) G c will be forwarded to the classifierf c for predicting the labels of the original graph.
Although Theorem E.4 assumes s c is known, inreal applications we do not know the specific s c .
Hence, in experiments, we select s c according tothe validation performance.
To thoroughly studythe effects of I ( (cid:98) G s ; Y ) comparing to CIGAv1,we stick to using the same s c and sampling process for CIGAv2, while CIGAv2 essentiallyrequires less specific knowledge about groundtruth r c hence achieving better empirical performance.
Moreover, once the sampled edgesare determined, the classifier GNN can take either the original feature of the input graph orthe learned feature from the featurizer as thenew node attributes for (cid:98) G c .
We select the architecture according to the validation performancefrom some random runs.
For the implementation of the information theoretic objectives, we will use CIGAv2 for elaboration while the implementation of CIGAv1 canbe obtained via removing the third term from CIGA v2.
Recall that CIGAv2 has the following formulation: [Equation] G [Equation]
= G − g
[Equation]
G
[Equation]
[Equation]
where (cid:98) G c
= g ( G ) , (cid:101) G c = g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G and G have the same label.
In Sec. 3.3, weintroduce a contrastive approximation for I ( (cid:98) G c ; (cid:101) G c | Y ) : I ( (cid:98) G c ; (cid:101) G c | Y ) ≈ [Equation] [Equation]
G c , (cid:101) G
[Equation]
log [Equation] 36) where positive samples ( (cid:98) G c , (cid:101) G c ) are the extracted subgraphs of graphs that have the same label of
G , negative samples are those with different labels, P g ( G |Y =
Y ) is the pushforward distributionof P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label
Y , h (cid:98) G c , h (cid:101) G c , h G ic are the graph presentations of the estimated subgraphs, and ϕ is the similaritymetric for the graph presentations.
As M → ∞ , Eq. 36 approximates I ( (cid:98) G c ; (cid:101) G c | Y ) which canbe regarded as a non-parameteric resubstitution entropy estimator via the von Mises-Fisher kerneldensity [1, 41, 100].
While for the third term I ( (cid:98) G s ; Y ) and the constraint
I ( (cid:98) G s ; Y ) ≤
I ( (cid:98) G c ; Y ) , a straightforwardimplementation is to imitate the hinge loss:
[Equation]
(37) where N is the number of samples, I is a indicator function that outputs 1 when the interior conditionis satisfied otherwise 0 , and R (cid:98) G s and R (cid:98) G c are the empirical risk vector of the predictions for each sample based on (cid:98) G s and (cid:98) G c respectively.
One can also formulate Eq.
35 from game-theoretic perspective [14].
Finally, we can derive the specific loss for the optimization of CIGAv2 combining Eq.
36 and Eq. 37:
[Equation]
) where R (cid:98) G c , R (cid:98) G s are the empirical risk when using (cid:98) G c , (cid:98) G s to predict Y through the classifier.
Typically, we use a additional MLP downstream classifier
ρ s for (cid:98) G s in the classifier GNN.
h (cid:98) G c is thegraph representation of (cid:98) G c which can be induced from the GNN encoder either in the featurizer or inthe classifier.
α, β are the weights for I ( (cid:98) G c ; (cid:101) G c | Y ) and I ( (cid:98) G s ; Y ) , and ϕ is implemented as cosinesimilarity.
The optimization loss for CIGAv1 merely contains the first two terms in Eq.
The detailed algorithm for CIGA is given in the Algorithm 1, assuming the h (cid:98) G c is obtained via thegraph encoder in f c . Fig. F also shows a illustration of the working procedure of CIGA.
G Detailed Experimental Settings
In this section, we provide more details about our experimental settings in Sec. 4, including the datasetpreparation, dataset statistics, implementations of baselines, selection of models and hyperparametersas well as evaluation protocols.
G.1 Details about the datasets
We provide more details about the motivation and construction method of the datasets that are used inour experiments.
Statistics of the datasets are presented in Table 4.
SPMotif datasets.
We construct 3-class synthetic datasets based on BAMotif [115, 57] following [103], where the model needs to tell which one of three motifs (House, Cycle, Crane) that thegraph contains.
For each dataset, we generate 3000 graphs for each class at the training set, 1000graphs for each class at the validation set and testing set, respectively.
During the construction, wemerely inject the distribution shifts in the training data while keep the testing data and validationdata without the biases.
For structure-level shifts ( SPMotif-Struc ), we introduce the bias based [Table] [Table] on FIIF, where the motif and one of the three base graphs (Tree, Ladder, Wheel) are artificially(spuriously) correlated with a probability of various biases, and equally correlated with the othertwo.
Specifically, given a predefined bias b , the probability of a specific motif (e.g., House) and aspecific base graph (Tree) will co-occur is b while for the others is (1 − b ) / 2 (e.g., House-Ladder,
House-Wheel).
We use random node features for SPMotif-Struc, in order to study the influencesof structure level shifts.
Moreover, to simulate more realistic scenarios where both structure leveland topology level have distribution shifts, we also construct SPMotif-Mixed for mixed distributionshifts.
We additionally introduced FIIF attribute-level shifts based on SPMotif-Struc, where all of thenode features are spuriously correlated with a probability of various biases by setting to the samenumber of corresponding labels.
Specifically, given a predefined bias b , the probability that all of thenode features of a graph has label y
(e.g., y = 0 ) being set to y (e.g., X = 0 ) is b while for the othersis (1 − b ) / 2 (e.g., P ( X = 1 )
=
P ( X = 2 ) =
(1 − b ) / 2 ).
More complex distribution shift mixescan be studied following our construction approach, which we will leave for future works.
TU datasets.
To study the effects of graph sizes shifts, we follow Yehudai et al.
[112], Bevilacquaet al.
[11] to study the OOD generalization abilities of various methods on four of TU datasets [66],i.e., PROTEINS, DD, NCI1, NCI109 .
Specifically, we use the data splits generated by Yehudai et al.[112] and use the Matthews correlation coefficient as evaluation metric following [11] due to theclass imbalance in the splits.
The splits are generated as follows:
Graphs with sizes smaller than the50 -th percentile are assigned to training, while graphs with sizes larger than the 90 -th percentile areassigned to test.
A validation set for hyperparameters tuning consists of 10% held out examples fromtraining.
We also provide a detailed statistics about these datasets in table 5.
Graph-SST datasets.
Inspired by the data splits generation for studying distribution shifts on graphsizes, we split the data curated from sentiment graph data [120], that converts sentiment sentenceclassification datasets SST5 and SST-Twitter [89, 26] into graphs, where node features are generatedusing BERT [25] and the edges are parsed by a Biaffine parser [32].
Our splits are created accordingto the averaged degrees of each graph.
Specifically, we assign the graphs as follows:
Those thathave smaller or equal than 50 -th percentile averaged degree are assigned into training, those that have averaged degree large than 50 -th percentile while smaller than 80 -th percentile are assigned tovalidation set, and the left are assigned to test set.
For SST5 we follow the above process while for
Twitter we conduct the above split in an inversed order to study the OOD generalization ability of
GNNs trained on large degree graphs to small degree graphs.
CMNIST-sp.
To study the effects of PIIF shifts, we select the ColoredMnist dataset created in IRM [4].
We convert the ColoredMnist into graphs using super pixel algorithm introduced by Knyazev et al.[46].
Specifically, the original Mnist dataset are assigned to binary labels where images with digits0 − 4 are assigned to y = 0 and those with digits 5 − 9 are assigned to y = 1 .
Then, y will be flippedwith a probability of 0 . 25 .
Thirdly, green and red colors will be respectively assigned to imageswith labels 0 and 1 an averaged probability of 0 .
15 (since we do not have environment splits) for thetraining data.
While for the validation and testing data the probability is flipped to 0 .
DrugOOD datasets.
To evaluate the OOD performance in realistic scenarios with realistic distribution shifts, we also include three datasets from DrugOOD benchmark.
DrugOOD is a systematic OOD benchmark for AI-aided drug discovery, focusing on the task of drug target bindingaffinity prediction for both macromolecule (protein target) and small-molecule (drug compound).
The molecule data and the notations are curated from realistic ChEMBL database [62].
Complicated distribution shifts can happen on different assays, scaffolds and molecule sizes.
In particular, we select DrugOOD-lbap-core-ic50-assay , DrugOOD-lbap-core-ic50-scaffold , and
DrugOOD-lbap-core-ic50-size , from the task of Ligand Based Affinity Prediction which usesic50 measurement type and contains core level annotation noises.
For more details, we referinterested readers to Ji et al.
G.2 Training and Optimization in Experiments
During the experiments, we do not tune the hyperparameters exhaustively while following thecommon recipes for optimizing GNNs.
Details are as follows.
GNN encoder.
For fair comparison, we use the same GNN architecture as graph encoders for allmethods.
By default, we use 3 -layer GNN with Batch Normalization [39] between layers and JKresidual connections at last layer [105].
For the architectures we use the GCN with mean readout [45]for all datasets except Proteins where we empirically observe better validation performance with a
GIN and max readout [106], and for DrugOOD datasets where we follow the backbone used in thepaper [40], i.e., 4 -layer GIN with sum readout.
The hidden dimensions are fixed as 32 for SPMotif,
TU datasets, CMNIST-sp, and 128 for SST5, Twitter and DrugOOD datasets.
Optimization and model selection.
By default, we use Adam optimizer [43] with a learning rate
of1 e − 3 and a batch size of 32 for all models at all datasets.
Except for DrugOOD datasets, we use abatch size of 128 following the original paper [40].
To avoid underfitting, we pretrain models for 20epochs for all datasets, except for CMNIST and Twitter where we pretrain 5 epochs and for SST5we pretrain 10 epochs, because of the dataset size and the difficulty of the task.
To avoid overfitting,we also employ an early stopping of 5 epochs according to the validation performance.
Meanwhile,dropout [90] is also adopted for some datasets.
Specifically, we use a dropout rate of 0 . 5 for CMNIST,
SST5, Twitter, DrugOOD-Assay and DurgOOD-Scaffold, 0 .
1 for DrugOOD-Size according to thevalidation performance, and 0 .
3 for TU datasets following the practice of Bevilacqua et al.
Implementations of baselines.
For implementations of the interpretable GNNs, we use the authorreleased codes [119, 77], where we use the codes provided by the authors 5 for DIR c[103] whichis the same as the author released codes.
During the implementation, we use the same s c for allinterpretable GNN baselines, chosen from { 0 . 1 , 0 . 2 , 0 . 25 , 0 . 3 , 0 . 4 , 0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 .
9 } accordingto the validation performances, and set to 0 .
25 for SPMotif following Wu et al.
3 for Proteinsand DD, 0 .
6 for NCI1, 0 .
7 for NCI109, 0 .
8 for CMNIST-sp, 0 . 5 for SST5 and Twitter, and 0 .
8 for
DrugOOD datasets, respectively.
Empirically, we observe that the optimization process in GIB canbe unstable during its nested optimization for approximating the mutual information of the predictedsubgraph and the input graph.
We use a larger batch size of 128 or reduce the nested optimization stepsto be lower than 20 for stabilizing the performance.
If the optimization failed due to the instabilityduring training, we will select the results with best validation accuracy as the final outcomes.
Although
SPMotif-Struc is also evaluated in DIR, we find the results are inconsistent to the results reported bythe author, because DIR adopts Last Epoch Model Selection which is different from the claim that they select models according to the validation performance , i.e., line 264 to line 278in train/spmotif dir.py from the commit 4b975f9b3962e7820d8449eb4abbb4cc30c1025dof https://github.com/Wuyxin/DIR-GNN .
We select the hyperparamter for the proposed DIRregularization from { 0 . 01 , 0 . 1 , 1 , 10 } according to the validation performances at the datasets, whilewe stick to the authors claimed hyperparameters for the datasets they also experimented with.
For invariant learning, we refer to the implementations in DomainBed [34] for IRM [4], V-Rex [49]and IB-IRM [2].
Since the environment information is not available, we perform random partitions onthe training data to obtain two equally large environments for these objectives.
Moreover, we selectthe weights for the corresponding regularization from { 0 . 01 , 0 . 1 , 1 , 10 , 100 } for these objectivesaccording to the validation performances of IRM and stick to it for others, since we empiricallyobserve that they perform similarly with respect to the regularization weight choice.
For EIIL [23],we use the author released implementations about assigning different samples the weights for beingput in each environment and calculating the IRM loss.
Besides, for CNC [122], we follow the algorithm description to modify the sampling strategy insupervised contrastive loss [42] based on a pretrained GNN optimized with ERM, and choose theweight for contrastive loss using the same grid search as for CIGA.
Implementations of CIGA.
For fair comparison, CIGA uses the same GNN architecture for GNNencoders as the baseline methods.
We did not do exhaustive hyperparameters tuning for the loss
Eq.
By default, we fix the temperature to be 1 in the contrastive loss, and merely search αfrom { 0 . 5 , 1 , 2 , 4 , 8 , 16 , 32 } and β from { 0 . 5 , 1 , 2 , 4 } according to the validation performances.
For
CMNIST-sp, we find larger β are required to get rid of intense spurious node features hence weexpand the search range for β to { 0 . 5 , 1 , 2 , 4 , 16 , 32 } , For Graph-SST datasets, we search α from{ 0 . 5 , 1 , 2 , 4 } as we empirically find that increasing α does not help increase the performance withfew random runs.
Besides, we also have various implementation options for obtaining the featuresin (cid:98) G c , for obtaining h (cid:98) G c , as well as for obtaining predictions based on (cid:98) G s .
By default, we feed thegraph representations of featurizer GNN to the classifier GNN, as well as to the contrastive loss.
Forclassifying G based on (cid:98) G s , we use a separate MLP downstream classifier in the classifier GNN f c .
The only exception is for the CMNIST-sp dataset where the spurious correlation is stronger thanthe invariant signal.
Directly feeding the graph representations from the featurizer GNN can easilyoverfit to the shortcuts hence we instead feed the original features to the downstream classifier GNN.
There can be more other options, such as using separate graph convolutions on (cid:98) G s or (cid:98) G c , which weleave for future work.
Evaluation protocol.
We run each experiment 10 on TU datasets and 5 times for others where therandom seeds start from 1 to the number of total repeated times.
During each run, we select themodel according to the validation performance and report the mean and standard deviation of thecorresponding metrics.
G.3 Software and Hardware
We implement our methods with PyTorch [72] and PyTorch Geometric [29].
We ran our experimentson Linux Servers with 40 cores Intel(R)
Xeon(R)
Silver 4114 CPU @
2.20GHz, 256 GB Memory,and Ubuntu 18.04 LTS installed.
GPU environments are varied from 4 NVIDIA RTX 2080Ti graphicscards with CUDA 10.2, 2 NVIDIA RTX 2080Ti and 2 NVIDIA RTX 3090Ti graphics cards with
CUDA 11.3, and NVIDIA TITAN series with CUDA 11.3.
G.4 Additional Analysis
Hyperparameter sensitivity analysis.
To examine how sensitive CIGA is to the hyperparamtersα and β for contrastive loss and hinge loss, respectively, under different distribution shifts.
Weconduct experiments based on the hardest datasets from each table (i.e., SPMotif-Mixed with thebias of 0 .
9 , DrugOOD-Scaffold and the NCI109 datasets from Table 1, Table 2, and Table 3,respectively.)
To increase the difficulty, we search for more fine-grained spaces for both parameters,i.e., { 0 . 1 , 0 . 5 , 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 } .
During changing the value of β , we will fix the α to a specificvalue under which the model has a relatively good performance (but not the best, to fully examine therobustness of CIGA in practice).
During the sensitivity tests, we follow the evaluation protocol asthat used for the main experiments.
The results are shown in Fig.
14 and Fig. 15.
[Figure]
c
[Figure]
[Figure]
From the results above
[Figure]
we can see that both CIGAv1 and CIGAv2 are robust to different valuesof α and β , respectively, across different datasets and distribution shifts.
Notably, in Fig.
14, whenthe coefficient α for the contrastive loss become too small, the invariance of the identified invariantsubgraphs (cid:98) G c may not be guaranteed, resulting worse performances.
Moreover, when α becomestoo large, it may affect the optimization and yield worse performances.
In SPMotif datasets, theworse performances can be observed via the large variances as well.
Similarly for β , as shown in
Fig.
15, when β becomes too small, some part from the spurious subgraph may still be contained inthe estimated invariant subgraphs.
While if β becomes too large, there might be part of (cid:98) G c beingeliminated.
Although both CIGAv1 and CIGAv2 are robust to the changes of α and β , the intrinsicdifficult optimization in OOD generalization algorithms including the proposed CIGA in our work,still require a more proper and smooth optimization process [18].
[Table]
Running time analysis.
To examine how much computational overhead is induced by the architectureand the additional objectives in CIGA, we analyze and compare the averaged training time of differentmethods on DrugOOD-Scaffold.
Factors that could affect the running time such as GNN backbone,batch size, and the running devices (NVIDIA RTX 2080Ti, Linux Servers with 40 cores Intel(R)
Xeon(R)
Silver 4114 CPU @
2.20GHz, 256 GB Memory, and Ubuntu 18.04 LTS), are fixed the sameduring the testing.
The results are shown as in Table.
It can be found that CIGA is the only OODmethod that outperforms ERM by a non-trivial margin with a relatively low additional computationaloverhead.
Single environment OOD generalization.
The theory of invariant learning fundamentally assumethe presence of multiple environments [75, 4].
However in practice, it does not always hold, whichwould inevitably fail all of the invariant learning solutions [4, 49, 23, 2], including CIGA.
[Table]
[Table]
Nevertheless, to examine how CIGA performs under various realistic scenarios, we conduct anadditional experiment based on DrugOOD-Assay.
We select samples that are from the largest assaygroup (i.e., the biochemical functionalities of these molecules are tested and reported under the sameexperimental setup in the lab)
The results are separated and shown in Table 7 and Table 8.
Besides the baselines, we also show the “Oracle” performances from the main table, to demonstratethe performance gaps.
From the Table 7 and Table 8, we can see that, both CIGAv1 and CIGAv2 maintain their state-of-theart performances even in the single training environment setting.
We hypothesize that enforcing themutual information between the estimated (cid:98) G c also helps to retain the invariance even under the singletraining environment setting.
That may partially explain why CNC can bring some improvements.
We believe it is an interesting and promising future directions to develop in-depth understanding andbetter solutions under this circumstance.
[Figure]
G.5 Interpretation Visualization
[Figure]
Since we use the interpretable GNN architecture to implement CIGA 6 , it brings an additional benefitthat provides certain interpretation for the predictions automatically, which may facilitate humanunderstanding in practice.
First, we provide some interpretation visualizations in SPMotif-Struc and SPMotif-Mixed datasets,under the biases of 0 .
6 and 0 . 9 .
Shown in Fig.
16 to Fig.
19, we use pink to color the ground truthnodes in G c , and denote the relative attention strength with edge color intensities.
Besides, we also provide some interpretation visualization examples in DrugOOD datasets.
Shown in
Fig.
20 to Fig.
25, we use the edge color intensities to denote the attentions of models that pay tothe corresponding edge.
Some interesting patterns can be found in the molecules shared with thesame label, which could provide insights to the domain experts when developing new drugs.
Webelieve that, because of its superior OOD generalization performance on graphs, CIGA can havehigh potential to push forward the developments of AI-Assisted Drug Discovery, and enrich the AItools for facilitating the fundamental practice of science in the future.
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
C
[Figure]
C O
C C C
C C
Cl
C C
[Figure]
N
[Figure]
O C C C
C
C C C C Br
C
S C
N
C C
C Br
C N C
N drugood_lbap_core_ic50_size: y= (a)
C C O C O C
C C C C
C O
C C
C C C
C C O
C C drugood_lbap_core_ic50_size: y= (b)
C O C O C C
C C
N C O C C
C C
C C
C C C O
C C
[Figure]
: y
[Figure]
(c)