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 [44, 34, 95, 102, 103].
However, it assumes that the training andtest graphs are drawn from the same distribution, which is often violated in reality [36, 46, 37, 39].
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, 23].
Such out-of-distribution (OOD) generalization failuresbecome the major roadblock for practical applications of graph representation learning [39].
Meanwhile, enabling OOD generalization on regular Euclidean data has received surging attentionand several solutions were proposed [4, 78, 10, 48, 22, 47, 2].
In particular, the invariance principlefrom causality is at the heart of those works [73, 71, 76].
The principle leverages the Independent
Causal Mechanism (ICM) assumption [71, 74] and implies that, model predictions that only focus onthe causes of the label can stay invariant to a large class of distribution shifts [73, 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 [109, 11, 98].
Onthe other hand, each of the shifts can spuriously correlate with labels in different modes [4, 68, 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[22]
Yes
R
No
PIIF DANN[30]
N/A
R
Yes
N/A MatchDG[59]
N/A
R
Yes
FIIF GroupDro[78]
N/A
R
Yes
N/A CNC[119]
N/A
R
No
N/A GIB[116]
Yes
G
No
FIIF DIR[100]
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 [36].
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) [71] 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 [39], 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, 22, 2], Group Distributionally
Robust Optimization [48, 78, 119], Domain Adaption and Domain Generalization [30, 89, 51, 26, 59,96] 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, 22], or under specific SCM assumptions [4], makingthe theoretical results hardly able to generalize to the complicated graph data [77] that can havemultiple types of distribution shifts [68].
Group Distributionally Robust Optimization that minimizesthe gap between worst group risk and average risk [48, 78, 119], and Domain Adaption/Generalizationmethods that aim to learn class-conditional domain invariant representations [30, 89, 51, 26, 96], cannot guarantee a min-max optimal predictor without additional assumptions [121, 4, 2].
Moreover,most existing methods require environment labels that are however expensive to obtain in graphs,which limits their applications to graphs [4, 48, 2, 78, 30, 89, 26, 59].
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 [112, 117].
Although some may leverage causality to justifythe generated explanation [52], 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 [116]and causality [100], respectively.
However, they focus on graphs and shifts generated under a specific
SCM.
Although one of them can provide theoretical guarantee for OOD generalization [116] 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.
[99] propose OOD generalization algorithms on graphs for the task ofnode 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 [93, 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 [71,2].
Thus, we will first formulate the data generation process with structural causal model and latentvariable model [71, 74, 49], 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 isgenerated 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 [49, 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, 120, 49].
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 [74].
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 notbe affected by the interventions (or changes) of E [71, 74].
For example, certain properties of amolecule can usually be described by a sub-molecule, or a functional group, which is invariantacross different species, or assays [12, 88, 39].
On the contrary, the generation of G s and G willbe affected by a subset of environment latent variable E G ⊆ E .
Thus, graphs collected fromdifferent environments (or domains) can have different distributions of structure-level properties(e.g., graph sizes [11, 98]) as well as feature-level properties (e.g., homophily [60, 17, 98]).
In fact, Assumption 2.1 is compatible to many graph generation models by specifying the function classes off G c gen , f G s gen and f G gen [85, 55, 113, 57].
Since our goal is to characterize the potential distribution shiftsin Assumption 2.1, we focus on building a general SCM that is compatible to many graph familiesand leave graph family specifications and their implications to OOD generalization in future works.
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 [71, 74].
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, 68].
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, 83, 84],which also serves as the necessary separation assumption for a classification task [66, 16, 62].
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.
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 [56, 100] 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 [93] 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
ERM tends to overfit to the shortcuts or spurious correlations presented in specific substructures or
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 ) [77].
In addition to IRM, the failure would also happen foralternative objectives [48, 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, 68].
Moreover, non-trivial environment partitions or labels are required for performanceguarantee of these OOD objectives [4, 48, 78, 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 [64, 27, 36].
Alternative options such as randompartitions tend not to alleviate the issue [22, 53], 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 [74], 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 [71].
Consequently, for a GNN with a permutationinvariant encoder h : G → R h and a downstream classifier ρ : R h
→ Y , if h can recover the information of C from G in the learned graph representations, then the learning of ρ resemblestraditional ERM [93] 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 [104, 106].
Specifically, we realize the alignment by decomposing a GNNinto two sub-components 2 : 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]
where R ( f c ( ˆ G c )) is the empirical risk of f c that takes ˆ G c as inputs to predict label Y for G , and ˆ G cis the intermediate subgraph containing information about C and hence needs to be independent of E .
Moreover, the extracted G c can either share the same graph space with input G or has its own spacewith latent node and edge features, depending on the specific graph generation process.
In practice,architectures from the literature of interpretable GNNs are compatible with CIGA [117], hence canserve as practical 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 can not satisfy the condition ˆ G c ⊥⊥ E .
Formally, merely minimizing R ( f c ( ˆ G c ))is equivalent to maximizing a variational lower bound of I ( ˆ G c ; Y ) [3, 116], which may include asubgraph from G s in ˆ 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 objectives [4, 48, 2, 78, 30, 89, 26, 59], making the identification of G c more challenging.
Optimization objective.
To mitigate this issue, we need to translate other properties of G c from SCMsin Sec. 2.2 into differentiable and equivalent conditions to the independence constraint ˆ G c ⊥⊥ E .
To start, consider a simplistic setting where all the invariant subgraphs G c have the same size s c ,i.e., | G c | =
s c 3 .
When maximizing I ( ˆ G c ; Y ) by minimizing R ( f c ( ˆ G c )) in Eq.
1, we are trying toextract all of the informative parts in G about Y into ˆ G c .
For FIIF (Fig. 2(b)), as G c already containsthe maximal possible information in G about Y , G c is a solution to max I ( ˆ G c ; Y ) .
However, somesubgraph of G c can be replaced by some subgraph of G s that is equally informative about Y .
For
PIIF (Fig. 2(c)), there also exists some subgraph of G s that contains additional information about Ythan G c , hence ˆ G c is more likely to involve some subgraph of G s .
Thus, the new condition needs toeliminate the auxiliary subgraphs of ˆ G c from G s such that the estimated ˆ G c can only contain G c .
Recall that for both FIIF and PIIF SCMs (Fig. 2), given two environments e 1 and e 2 : if G c appears inboth e 1 and e 2 , the correctly identified ˆ G e 1 c , ˆ G e 2 c in e 1 and e 2 tend to have higher mutual informationabout the other, i.e., ( G c , G c ) ∈ arg max I ( ˆ G e 1 c ; ˆ G e 2 c ) .
While for G c and another G c ′ correspondingto a different C ′ ̸
= C , if they appear in the same environment, then including any subgraph from G s in the estimated ˆ G c , ˆ G c ′ for G c , G c ′ will enlarge their mutual information, i.e., ( G c , G c ′ ) ∈arg
min I ( ˆ G c ; ˆ G c ′ ) .
Thus, we can derive another important property about G c :
[Equation]
where ˆ G c and ˜ G c share the same C while ˆ G c ′ corresponds to a different c ′ .
In practice, we are alsonot given C .
However, since C and Y shares a stable causal relationship in both FIIF and PIIF SCMs,
Y can serve as an alternative of C in Eq. 2.
Moreover, when both I ( ˆ G c ; ˜ G c | C = c ) and I ( ˆ G c ; Y )are maximized, I ( ˆ G c ; ˆ G c ′ | C
= c ′ , c ′ ̸
= c ) is automatically minimized, otherwise all classes willcollapse to trivial solutions which is not possible given I ( ˆ G c ; Y ) being maximized.
Thus, we canreplace the independence condition in Eq. and obtain the following objective (CIGAv1):
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
In Theorem 3.1,we discuss why Eq. 3 is equivalent to Eq.
Although being effective, Eq. 3 requires a strongassumption about the size of G c .
However, the size of G c is usually unknown or changes for different C s. In this circumstance, maximizing Eq.
3 without additional constraints would lead to the presenceof subgraphs of G s in ˆ G c .
For instance, ˆ G c =
G is a trivial solution to Eq.
3 when s c = ∞ .
To circumvent this limitation, we further resort to the properties of G s .
In both FIIF and PIIF SCMs (Fig. 2), G s and G c share certain overlapped information about Y via some subgraphs.
Whenmaximizing I ( ˆ G c ; ˜ G c | Y ) and I ( ˆ G c ; Y ) , the appearances of these subgraphs in ˆ G c will not affect theoptimality.
On the other hand, it can reduce the mutual information between the left part ˆ G s =
G − ˆ G cand Y , i.e., I ( ˆ G s ; Y ) .
In other words, by maximizing I ( ˆ G s ; Y ) , we can avoid involving additionalsubgraphs from G s into ˆ G c .
Meanwhile, to avoid trivial solution that G c ⊆ ˆ G s during maximizing
I ( ˆ G s ; Y ) , we can leverage the better clustering property of G c implied by Assumption 2.4 to derivethe constraint
I ( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) .
Thus, we can obtain a new objective CIGAv2 as follows:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
We also provethe 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
(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 ( ˆ G c ; ˜ G c | Y ) could be highly expensive [92, 8].
However, contrastive learning withsupervised sampling provides a practical solution for the approximation [41, 19, 79, 92, 8]:
[Equation]
where positive samples ( ˆ G c , ˜ 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 push-forward distributionof P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label Y , 282 P ( G |Y
̸ = Y ) refers to the distribution of G given the label that is different from Y , h ˆ G c , h ˜ 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 ( ˆ G c ; ˜ G c | Y ) , which can be regarded as a nonparameteric resubstitution entropy estimator via the von Mises-Fisher kernel density [1, 40, 97].
Thus,plugging it into Eq. 3 and Eq.
4 can relieve the issue of approximating I ( ˆ G c ; ˜ G c | Y ) in practice.
[Table]
For the implementation of I ( ˆ G s ; Y ) and the constraint I ( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) in CIGAv2, a practicalchoice is to follow the idea of hinge loss, I ( ˆ G s ; Y )
≈ 1 N R ˆ G s · I ( R ˆ G s ≤ R ˆ 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 ˆ G s and R ˆ G c are the empirical risk vector of the predictions for each sample basedon the corresponding ˆ G s and ˆ 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 [26, 59, 119], 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 [119], 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 [83, 84].
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 [75], GIB [116] or DIR [100] 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 [100] 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 [39] from AI-aided Drug Discovery with Assay,
Scaffold, and Size splits, convert the ColoredMNIST from IRM [4] using the algorithm from Knyazevet al.
[45] to inject attribute shifts, and split Graph-SST [117] to inject degree biases.
To compare withprevious specialized OOD methods for graph size shifts [109, 11], we use the datasets in Bevilacquaet al.
[11] that are converted from TU benchmarks [64].
More details can be found in Appendix G.1.
Baselines and our methods.
Besides the ERM, we also compare with SOTA interpretable GNNs,
GIB [116], ASAP Pooling [75], and DIR [100], to validate the effectiveness of the optimizationobjective in CIGA.
We use the same selection ratio (i.e., s c ) for all models.
On the other hand,to validate the effectiveness of the decomposition in CIGA, we compare CIGA with SOTA OODobjectives including IRM [4], v-Rex [48] and IB-IRM [2], for which we apply random environmentpartitions.
We also compare CIGA with EIIL [22] and CNC [119] that does not require environmentlabels, where CNC [119] 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 the DrugOOD datasetswhere we use ROC-AUC following [39], and for the TU datasets where we use Matthews correlation
[Table]
coefficient following [11]. We repeat the evaluation multiple times, select models based on thevalidation performances, and report the mean and standard deviation of the corresponding metric.
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 methodsconverge 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]
[Table]
OOD generalization performance 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 [33, 39].
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 tothem?
[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.
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 Review Board (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 amountspent on participant compensation?
[N/A] We didn’t conduct research with humansubjects.
Full Structural Causal Models on Graph
[-B More Discussions on Related Works and Future Directions 
 B.1 More backgrounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 Detailed related work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .-]
Generation {+
+}
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.
E
Theory
{+and Discussions 
+}
E.1 Challenges of OOD generalization on graphs. . . . . . . . . . . . . . . . . . . . .
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 . .
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
= ρ ( r 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) [103, 111, 67, 103,18, 65], 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 [93, 73, 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 [112, 117].
Although some may leveragecausality in explanation generation [52], 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 [116] and DIR [100], and both of themfocus on tackling FIIF spurious correlations (Assumption C.2).
The theoretical guarantees of GIBfollows the theory of information bottleneck [91], 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, 68].
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 [81, 82],physics [6, 80], and graph algorithms [90, 94, 104, 101].
Xu et al.
[106] study the neural networkextrapolation ability from a geometrical perspective.
Han et al.
[35] improve OOD drug discoveryby mitigating the overconfident misprediction issue.
Knyazev et al.
[45], Yehudai et al.
[109] 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.
[99] 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.
[99],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 [72, 83, 84].
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 [71, 74].
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 [73, 4].
Consequently, models built upon the invariantrepresentation can generalize to unseen environments or domains with guaranteed performance [73,76, 4, 78, 10, 47, 33, 48, 22, 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 [69] and generalization [118], 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, 68].
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.
[22] and Lin et al.
[53] 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.
[78] propose to regularize the worst groupin group distributionally robust optimization (GroupDro).
Zhang et al.
[119] 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[76, 20, 30, 89, 51, 26, 59, 96], 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 [33, 2], hence are not able to yield predictors with OOD generalizationguarantees.
We refer interested readers to Pearl [72], Sch ¨ olkopf [83], Sch ¨ olkopf et al.
[84] for anin-depth understanding, and Gulrajani and Lopez-Paz [33], 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 ) [122, 49].
Moreover, our results also partially explain the success of graph contrastive learning [114, 58, 115],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 [100] as they also focus on FIIF), which further implies Y ⊥⊥ S | ˆ G c .
Hence it is natural to use
Information Bottleneck (IB) objective [91] to solve for G c :
[Equation]
which explains the success of many existing works in finding predictive subgraph through IB [116].
However, the estimation of I ( ˆ 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 [99].
B.4 Discussions on limitations of CIGA and future directions
Better graph generative 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 [85, 55, 113, 57].
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 [107, 29, 54].
Thus, it is promising to leverage theprior knowledge for specific graph generative 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 [19, 79, 92, 114, 115].
A better augmentation or sampling strategy can benefit the
OOD generalization in general as shown by K ¨ ugelgen et al.
[49] 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 [41].
Thus, it is promising to leverage better augmentation and contrastivestrategy to improve the generalization ability in CIGA.
More sophisticated architectures.
The CIGA framework introduced in Sec.
3 can have multipleimplementations.
We choose interpretable architectures in our experiments for the purpose ofconcept verification.
More sophisticated architectures can be incorporated, such as identifying anddisentangling G c at the latent space [83, 84].
C Full Structural Causal Models on Graph Generation
Supplementary to the graph generation process in Sec. 2.2, we provide full SCMs on the graphgeneration 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 [49, 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, 120, 49].
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 (Scaffold), experimental environment for examining the chemicalproperty (Assay) [39].
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]. S C E G Z cX Z cA
Z sX
Z sA G c G s G (a) Graph Generation SCM
E
S Y
G
C (b) FIIF SCM
E
S Y
G
C (c) PIIF SCM
E
S
C
S 2 Y
G (d) MIIF SCM
Assumption C.1 (Graph Generation SCM) .
[Equation]
Specifically, the graph generation process is shown as Fig. 4(a).
The generative 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 [85, 55, 113, 57, 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 [85, 55, 113, 57], and leave specific graph generation processesand their implications to OOD generalization to future work, as discussed in Appendix B.4. Moreover, a subset of environment latent variable E G ⊆
E will affect the generation of G and
G s . Thus, graphs collected from different environments can have different structure-level propertiessuch as degrees, graph sizes, and subgraph densities, as well as feature-level properties such ashomophily [45, 109, 11, 17].
Meanwhile, all of them can spuriously correlated with the labelsdepending on how the underlying latent variables are interacted with each others.
The interactiontypes can be further divided into two axiom types FIIF and PIIF, as well as the mixed one MIIF.
Previous OOD methods such as GIB [116] and DIR [100] mainly focus on FIIF case, while otherssuch as IRM [4] mainly focuses on the PIIF case.
Evidences show that failing to model either of themwhen developing the OOD objectives can have serious performance degenerations in practice [5, 68].
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]
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 [71, 74].
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 2.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, 68].
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, 83, 84], which alsoserves as the necessary separation assumption for a classification task [66, 16, 62].
Assumption C.5 (Latent Separability) .
H ( C | Y ) ≤ H ( S | Y ) .
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.
[Figure]
GNNs are required to classify whether the graph contains a “house” or “cycle” motif, where the colors represent node
IRM fail when testing data not sufficiently by the data.
To begin with, we construct 3-class syntheticdatasets based on BAMotif [56] and follow Wuet al. [100] to inject spurious correlations between motif graph and base graph during thegeneration.
In this graph classification task, themodel needs to tell which motif the graph contains, e.g., “House” or “Cycle” motif, as shownin Fig.
We inject the distribution shifts inthe training data while keeping the test data andvalidation data without the biases.
For structurelevel shifts, we introduce the artificial bias basedon FIIF, where the motif and the base graph arespuriously correlated with a probability of various bias.
For mixed shifts, we additionally introduced attribute-level shifts based on FIIF, whereall of the node features are spuriously correlatedwith a probability of various bias.
The numberof training graphs is 600 for each class and thenumber of graphs in validation and test set is200 for each class.
More construction detailsare given in Appendix G.
For the GNN encoders, by default, we use 3layer GCN [44] with mean readout, a hiddendimension of 64 , and JK jump connections [102]at the last layer.
During training, we use a batchsize of 32 , learning rate of 1 e − 3 with Adamoptimizer [42], and batch normalization betweenhidden layers [38].
Meanwhile, to stabilize thetraining, we also use dropout [87] of 0 .
1 andearly stop the training when the validation accuracy does not increase till 5 epoch after first 20 epochs.
All of the experiments are repeated 5 times, and the mean accuracy as well as variance are reportedand plotted.
When using IRM objective [4], as the environment partitions are not available, wegenerate 2 environments with random partitions.
D.2 More discussions about failure case study in Sec. 2.3
In Fig. 6, 7, 8, 9, we investigate whether existing training objectives (ERM and IRM), adding moremessage passing, as well as using expressive GNNs, can improve the OOD generalization ability ongraphs.
Here we also provide a additional discussion in complementary to the discussions on OODgeneralization performance of ERM and IRM objectives in Sec. 2.3.
Can better architectures improve OOD generalization of GNNs?
Adding more message passing turns.
It is a common practice in GNNs to denoise the signalsby aggregating more neighbors with higher layers, or enhance the expressive power with morepowerful readout functions [102, 103, 108].
Aggregating neighbor information with more layers todenoise the input signal, or enhancing the expressivity with more powerful readout functions, aretwo common choices in GNNs to improve the generalization ability [102, 50, 103, 108].
However,in the experiments next, we empirically found that GCNs with more layers and more powerfulreadout operations are still sensitive to distribution shifts.
In particular, stacking more layers helpsdenoising certain shifts, while the OOD performance would drop more sharply when the biasincreases.
Intuitively, if the spurious features from nodes cannot be eliminated by the denoisingproperty of a deeper GNN, they would spread among the whole graph more widely, which in turnleads to stronger spurious correlations.
Besides, the spurious correlations would be more difficult tobe disentangled if there are distribution shifts at both structure-level and attribute-level.
Since thenode representations from hidden layers can also encode graph topology features [103], distributionshifts introduced through Z sA and Z sX will doubly mix at the learned features.
In the worst case, theinformation about Z cA and Z cX could be partially covered by or even replaced by Z sA and Z sX .
Thiswill make OOD generalization of message passing GNNs trained through ERM much more difficultor even impossible.
Besides, as the node representations of 1 ≤ i ≤ k -th layer can also encode graphtopology features [103], which, if spuriously correlated with labels through Z sA and entangled withpart of invariant node features, i.e., Z cX , in the worst case, can greatly improve the difficulty or evenmake the OOD generalization impossible 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 [103].
Afterthat, many follow-up variants are proposed to improve the expressivity of GNNs [65].
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, 23, 4, 78, 10, 48, 22, 47, 2].
Empirically, we use k -GNNs [63] to verify theintuition and observe similar failures for this provably more expressive GNN as basic GNN variants.
D.3 More empirical results about failure case study in Sec. 2.
[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.
[100] to construct dataset for Fig.
6, 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. 7, graph size shifts are added, which is exactly the SPMotif datasets used in
DIR [100].
Besides, in Fig.
8, 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.
6, 7, 8, 9,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 [32].
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.
6, when the spurious correlation grows stronger, the OOD performancebecome more unstable , e.g., higher variance.
The reason is that, GNNs sometimes can directly learn 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.
6, 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.[45] 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.
6 to Fig.
8, as well as Fig.
7 to Fig.
Observation IV: Using more powerful architectures can not improve the OOD performance.
From the sub-figures (b) and (c) in Fig.
6, 7, 8, 9, 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 [105].
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.
6, 7, 8, 9 we use 1 - 2 - 3 -GNNfollowing the algorithm of k -GNNs which is provably more expressive than 2 -WL test [63].
Whenthere are no graph size shifts, k -GNNs will have higher performance at the beginning.
When thereare graph size shifts, k -GNNs will have a lower initial performance at the beginning.
Then, as thespurious strength grows, k -GNNs can suddenly become seriously unstable, though k -GNNs can havehigher 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.3
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 [93, 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 ⊆ 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 [78, 48, 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.
[77], 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 [22] 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 [100], 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.
[100])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 [116] that focuses on discovering a informative subgraph for explanation,essentially can provide theoretical guarantees for FIIF spurious correlations.
Theoretically, (we copy 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 [100] as they also focuson FIIF), which further implies Y ⊥⊥ S | ˆ G c .
Hence it is natural to use Information Bottleneck (IB)objective [91] to solve for G c :
[Equation]
which explains the success of many existing works in finding predictive subgraph through IB [116].
However, the estimation of I ( ˆ 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.
E Theory and Discussions
In this section, we provide proofs for propositions and theorems mentioned in the main paper.
E.1 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.2 Proof for theorem 3.1 (i)
Theorem E.1 (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 ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
The proof of Theorem E.1 is essentially to show the estimated ˆ G c through Eq.
14 is the underlying
G c , then the maximizer of I ( ˆ G c ; Y ) in Eq.
14 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 ( ˆ G c ; Y ) and thesecond term I ( ˆ G c ; ˜ G c | Y ) to conclude the proof.
We begin by introducing the following lemma:
Lemma E.2.
Given the same conditions as Thm.
E.1, I ( ˆ G c ; Y ) is maximized if and only if
I ( ˆ G c ; Y | E = e ) is maximized, ∀ e ∈ E tr .
The proof for Lemma E.2 is straightforward, given the condition that samples from each trainingenvironment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr .
Obviously, ˆ G c =
G c is amaximizer of I ( ˆ 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 ˆ G c = G pc ∪ G ps and G pc = G c ∩ ˆ G c .
For FIIF (Assumption 4(b)), it can not happen otherwise, let G lc
= G c − G pc , then we have:
[Equation]
[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 ˆ G c , we will use the second term to eliminate it:
[Equation]
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ 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 ( ˆ G c ; ˜ G c | Y ) is essentially maximizes I ( ˆ G c , E
= ˆ e ; ˜ G c , E = ˜ [Equation]
| Y ) , ∀ ˆ e, ˜ e ∈ E tr .
Hence, we have:
[Equation]
We claim
Eq.
can eliminate any potential subsets in the estimated ˆ G c .
[Equation]
Otherwise, suppose there are some subsetsˆ G ps ⊆ ˆ G s and ˜ G ps ⊆ ˜ G s contained in the estimated ˆ G c , ˜ G c , where ˆ G s , ˜ G s be the corresponding underlying G s s for ˆ G c , ˜ G c .
Let ˆ G ∗ c and ˜ G ∗ cbe the ground truth invariant subgraph G c s of ˆ Gand ˜ G , ˆ G lc = ˆ G ∗ c − ˆ G c and ˜ G lc = ˜ G ∗ c − ˜ G c bethe l eft (un-estimated) subsets from corresponding ground truth G c s, and ˆ G pc
= ˆ G ∗ c − ˆ G lc and˜ G pc = ˜
G ∗ c − ˜
G lc be the complement, or equivalently, the p artial ˆ G ∗ c , ˜
G ∗ c that are estimated inˆ G c , ˜ G c , respectively.
We can also define similar counterparts for G s : ˆ G ps , ˜
G ps are the partial ˆ G s , ˜ G s s contained in the estimated ˆ G c , ˜ G cwhile ˆ G ls , ˜ G ls are the left subsets ˆ G s , ˜ G s , respectively.
Recall the constraint that | G c | =
s c , hence if ˆ G pc ⊆ ˆ G c , then a corresponding ˆ G lc = ˆ G ∗ c − ˆ G pc will bereplaced by ˆ G p s in ˆ G c .
In this case, we have:
[Equation]
where the second equality is due to E = ˆ e is determined so that H ( E = ˆ e | ˆ G c , Y ) = 0 .
Compared
Eq.
18 to that when ˆ G c
=
ˆ G ∗ c , we have the entropy change as:
[Equation]
Let ϵ
= H ( ˆ G ps | ˆ 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.
19, H ( ˆ G ps | ˆ G pc , E =
ˆ e, Y )
=
0while
H ( ˆ G lc | ˆ G pc , E = ˆ e, Y ) > {+
+} since ˆ G lc can not be determined when given ˆ G pc , E = ˆ e, Y .
Thus,when some subset from G s is included in ˆ G c , it will minimize H ( ˆ G c , E
= ˆ e | Y ) .
In the next, we willshow how ϵ =
H
( ˆ G ps | ˆ G pc , E = ˆ e, Y ) can be cancelled thus leading to a smaller H ( ˆ G c , E
= ˆ e | Y ) ,by considering the second term H ( ˆ G c , E
= ˆ e | ˜ G c , E = ˜
e, Y ) .
As for H ( ˆ G c , E = ˆ e | ˜ G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possiblecases into two:
(i) One of ˆ G c and ˜ G c contains some subset of G s , i.e., ˆ G c
contains some
ˆ G p s ⊆ ˆ G s ; (ii) Both ˆ G c and ˜ G c contain some ˆ G ps ⊆
ˆ G s and ˜ G ps ⊆ ˜ G s , respectively
For (i), we have:
[Equation]
Thus, H ( ˜ G ps | ˆ G c , E = ˜
e, Y, ˜
G lc , E = ˆ e ) =
and Thus, we have:
[Equation]
Combing ∆ H ( ˆ G c , E = ˆ e | Y ) , we have:∆
[Equation]
ˆ
[Equation]
c
[Equation]
where the last equality is because of the independence of ˆ G ps between ˜ G c , E =
˜
e conditioned on
Y, E = ˆ e .
Since conditioning will lower the entropy for both discrete and continuous variables [21,110], we have:
[Equation]
which implies the existence of ˆ G ps in ˆ G c will lower down the second term in Eq.
14 for the case (i).
For (ii), we have:
[Equation]
Similar to (i), H ( ˆ G ps | ˜ G pc , ˜ G ps , E = ˜
e, Y, ˆ G pc , E = ˆ e ) can be cancelled out with H ( ˆ G ps | ˆ G pc , E = ˆ e, Y ) .
Then, we have:
[Equation]
Since additionally conditioning on ˆ G ps in H ( ˆ G lc , E = ˆ e | ˜ G pc , ˜ G ps , E = ˜
e, Y ) can not lead to newinformation about ˆ G lc , we have:
[Equation]
which follows that ∆ I ( ˆ G c , E = ˆ e ; ˜ G c , E = ˜
e | Y ) < 0 .
To summarize, the ground truth G c is the only maximizer of the objective (Eq. 14), hence solving forthe objective (Eq. 14) can elicit an invariant GNN.
E.3 Proof for theorem 3.1 (ii)
Theorem E.3 (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]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
Similar to the proof for Theorem E.1, to prove Theorem E.3 is essentially to show the estimatedˆ G c through Eq.
27 is the underlying G c , hence the minimizer of Eq.
27 elicits an invariant GNNpredictor.
In the next, we also begin with a lemma:
Lemma E.4.
Given data generation process as Theorem E.3, 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.4.
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.4, we know ˆ G c at least contains some subset of the underlying G c , otherwise theconstraint I ( ˆ G s ; Y ) ≤ I ( ˆ G c ; Y ) will be violated since G c ⊆ ˆ G s in this case.
Assuming there are some subset of G s contained in ˆ G c , without loss of generality, we can divide allof the possible cases about ˆ G c into two:
(i) ˆ G c only contains a subset of the underlying G c ;(ii) ˆ G c contains a subset of the underlying G c as well as part of the underlying G s ;
[Equation]
Before the discussion, let us inherit the notationsof subsets of G c , G s from the proof for Theorem E.1: Let ˆ G ∗ c and ˜
G ∗ c be the ground truthinvariant subgraph G c s of ˆ G and ˜ G , ˆ G lc = ˆ G ∗ c −ˆ G c and ˜ G lc =
˜
G ∗ c − ˜ G c be the l eft (un-estimated subsets from ground truth G c s,and ˆ G pc = ˆ G ∗ c − ˆ G lc and ˜ G pc = ˜
G ∗ c − ˜ G be the complement, or equivalently the artial ˆ G ∗ c , ˜ G ∗ c are G c , ˜ G c , respectively.
Sim ilarly, ˆ G ps , ˜ G ps are the partial ˆ G s , ˜ G s s contained in
˜ G First of all, case (i) cannot hold because, whenmaximizing I ( ˆ G c ; ˜ G c | Y ) , if ∃ ˆ G lc = ˆ G ∗ c − ˆ G c ,as shown in the proof for Theorem E.1, including ˆ G lc into ˆ G c can always enlarge I ( ˆ G c ; ˜ G c | Y ) ,while not affecting the optimality of I ( ˆ G s ; Y )
+
I ( ˆ G c ; Y ) by re-distributing ˆ G lc from ˆ G s to ˆ G c .
Consequently, ˆ G ∗ c must be included in ˆ G c , i.e.,ˆ G ∗ c ⊆ ˆ G c .
As for case (ii), recall that, by the condition of equally distributed training samples from each trainingenvironment, maximizing I ( ˆ G c ; ˜ G c | Y ) is essentially maximizing I
( ˆ G c , E = ˆ e ; ˜ G c , E = ˜
e | Y ) ,∀ ˆ e, ˜ e ∈ E tr , hence, we have:
[Equation]
We claim Eq.
28 can eliminate any potential subsets in the estimated ˆ G c .
Similarly, we have:
[Equation]
where the second equality is due to E
= ˆ e is determined.
Compared to the case that ˆ G c
= ˆ G ∗ c , wehave:
[Equation]
Then, as for H ( ˆ G c , E = ˆ e | ˜ G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possiblecases into two:
(a) ˆ G c contains some ˆ G ps ⊆ ˆ G s ;(b) Both ˆ G c and ˜ G c contain some ˆ G ps ⊆ ˆ G s and ˜ G ps ⊆ ˜ G s , respectively.
For (a), we have:
[Equation]
[Equation]
Similarly to the proof for Theorem E.1, when considering ∆ I ( ˆ G c ; ˜ G c | Y ) , the effects of
H ( ˆ G ps | ˜ G c , E = ˜
e, Y, ˆ G ∗ c , E = ˆ e ) is cancelled out by H ( ˆ G ps | ˆ G ∗ c , E = ˆ e, Y ) .
Hence, we have:
[Equation]
For (b), we have:
[Equation]
Similarly, H ( ˆ G ps | ˜ G ∗ c , ˜ G ps , E = ˜
e, Y, ˆ G ∗ c , E = ˆ e ) = 0 can also be cancelled out by H ( ˆ G ps | ˆ G ∗ c , E =ˆ e, Y ) .
Moreover, for H ( ˆ G ∗ c | ˜ G ∗ c , ˜ G ps , E = ˜
e, Y, E
= ˆ e ) , ˜ G ps can not bring no additional informationabout ˆ G ∗ c , when conditioning on ˜ G ∗ c , Y, E = ˜
e .
Hence, we also have:
[Equation]
To summarize, when maximizing I ( ˆ G c ; ˜ G c | Y ) , including any ˆ G p s ⊆ ˆ G ∗ s can not bring additionalbenefit while affecting the optimality of I ( ˆ G s ; Y )
+ I ( ˆ G c ; Y ) .
More specifically, when consideringthe changes to I ( ˆ G s ; Y ) + I ( ˆ G c ; Y ) ,
∀ G ps ⊆
G s , we have
[Equation]
c −
[Equation]
≤ I
[Equation] [Equation] , while I ( Y ; ˆ G ∗ c , G ps ) =
I ( Y ; ˆ G ∗ c )
+ I
( Y ; ˆ G ps | ˆ G ∗ c ) , ∀ e ∈ E tr .
Consequently,
[Equation]
Hence, only the underlying G c is the solution to Eq. 27, which implies that solving for the objective(Eq. 27) 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 cand a classifier f c : G c → Y .
Specifically, for the implementation of Featurizer, we choose one of thecommon practices GAE [43] for calculating the sampled weights for each edge.
More formally, thesoft mask is predicted through the following equation:
[Equation]
If a sampling ratio s c is predetermined, we sample s c of total edges with the largest predicted weightsas a soft estimation of ˆ G c .
Then, the estimated ˆ G c will be forwarded to the classifier f c for predictingthe labels of the original graph.
Although Theorem E.1 assumes s c is known, in real applicationswe do not know the specific s c .
Hence, in experiments, we select s c according to the validationperformance.
To thoroughly study the effects of I ( ˆ G s ; Y ) comparing to CIGAv1, we stick to usingthe same s c and sampling process for CIGAv2, while CIGAv2 essentially requires less specificknowledge about ground truth r c hence achieving better empirical performance.
Moreover, once thesampled edges are determined, the classifier GNN can take either the original feature of the inputgraph or the learned feature from the featurizer as the new node attributes for ˆ G c .
We select thearchitecture according to the validation performance from some random runs.
For the implementation of the information theoretic objectives, we will use CIGAv2 for elaborationwhile the implementation of CIGAv1 can be obtained via removing the third term from CIGAv2.
Recall that CIGAv2 has the following formulation:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
In Sec. 3.3, weintroduce a contrastive approximation for I ( ˆ G c ; ˜ G c | Y ) :
[Equation]
where positive samples ( ˆ G c , ˜ 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 ˆ G c , h ˜ G c , h G ic are the graph presentations of the estimated subgraphs, and ϕ is the similarity metricfor the graph presentations.
As M → ∞ , Eq. 35 approximates I ( ˆ G c ; ˜ G c | Y ) which can be regarded asa non-parameteric resubstitution entropy estimator via the von Mises-Fisher kernel density [1, 40, 97].
While for the third term I ( ˆ G s ; Y ) and the constraint
I ( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) , a straightforwardimplementation is to imitate the hinge loss:
[Equation]
36) where N is the number of samples, I is a indicator function that outputs 1 when the interior conditionis satisfied otherwise 0 , and R ˆ G s and R ˆ G c are the empirical risk vector of the predictions for eachsample based on ˆ G s and ˆ G c respectively.
One can also formulate Eq.
34 from game-theoreticperspective [14].
Finally, we can derive the specific loss for the optimization of CIGAv2 combining Eq.
35 and Eq.
[Equation]
37) where R ˆ G c , R ˆ G s are the empirical risk when using ˆ G c , ˆ G s to predict Y through the classifier.
Typically, we use a additional MLP downstream classifier
ρ s for ˆ G s in the classifier GNN.
h ˆ G c is thegraph representation of ˆ G c which can be induced from the GNN encoder either in the featurizer or inthe classifier.
α, β are the weights for I ( ˆ G c ; ˜ G c | Y ) and I ( ˆ 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 ˆ G c is obtained via thegraph encoder in f c .
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.
samples samples constructed for
[Table]
SPMotif datasets.
We construct 3-class synthetic datasets based on BAMotif [112, 56] following [100], 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 basedon 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.
[109], Bevilacquaet al.
[11] to study the OOD generalization abilities of various methods on four of TU datasets [64],i.e., PROTEINS, DD, NCI1, NCI109 .
Specifically, we use the data splits generated by Yehudai et al. [Table] [109] 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 [117], that converts sentiment sentenceclassification datasets SST5 and SST-Twitter [86, 25] into graphs, where node features are generatedusing BERT [24] and the edges are parsed by a Biaffine parser [31].
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 thathave 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.[45].
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 [61].
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 [38] between layers and JKresidual connections at last layer [102].
For the architectures we use the GCN with mean readout [44]for all datasets except Proteins where we empirically observe better validation performance with a
GIN and max readout [103], and for DrugOOD datasets where we follow the backbone used in thepaper [39], 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 [42] 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 [39].
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 [87] 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 [116, 75], where we use the codes provided by the authors 4 for DIR c[100] 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 claimthat 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 [33] for IRM [4], V-Rex [48]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 [22],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 [119], we follow the algorithm description to modify the sampling strategy insupervised contrastive loss [41] 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. 37.
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 ˆ G c , for obtaining h ˆ G c , as well as for obtaining predictions based on ˆ 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 ˆ 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 easily overfit 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 ˆ G s or ˆ 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 [70] and PyTorch Geometric [28].
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.