Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs
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 are still limited.
Different from images, the complex nature of graphs poses unique challenges to adopting the invariance principle.
In particular, distribution shifts on graphs can appear in a variety of forms such as attributes and structures, making it difficult to identify the invariance.
Moreover, domain or environment partitions, which are often required by OOD methods on Euclidean data, could be highly expensive 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 causal models, concluding that OOD generalization on graphs is achievable when models focus only on subgraphs containing the most information about the causes of labels.
Accordingly, we propose an information-theoretic objective to extract the desired subgraphs that maximally preserve the invariant intra-class information.
Learning with these subgraphs is immune to distribution shifts.
Extensive experiments onsynthetic or real-world datasets, including a challenging setting – DrugOOD,from
AI-aided drug discovery, validate the superior OOD performance of CIGA 1 .
Introduction
Graph representation learning with graph neural networks (GNNs) has gained great success in tasks involving relational information [45, 35, 99, 106, 107].
However, it assumes that the training and test graphs are drawn from the same distribution, which is often violated in reality [37, 47, 38, 40].
The mismatch between training and test distributions, i.e., distribution shifts , introduced by some underlying environmental factors related to data collection or processing, could seriously degrade the performance of deployed models [7, 24].
Such out-of-distribution (OOD) generalization failures become the major roadblock for practical applications of graph representation learning [40].
Meanwhile, enabling OOD generalization on regular Euclidean data has received surging attention and several solutions were proposed [4, 81, 10, 49, 23, 48, 2].
In particular, the invariance principle from causality is at the heart of those works [76, 74, 79].
The principle leverages the Independent
Causal Mechanism (ICM) assumption [74, 77] and implies that, model predictions that only focus on the causes of the label can stay invariant to a large class of distribution shifts [76, 4, 2].
“Cycle” “House”
Classification Invariant Subgraph Identification
[Figure]
Algorithm
OODGuarantee Regime E Known SCMSupport
IRM[4]
Yes
R
Yes
PIIF IB-IRM[2]
Yes
R
Yes
PIIF&FIIF EIIL[23]
Yes
R
No
PIIF DANN[31]
N/A
R
Yes
N/A MatchDG[61]
N/A
R
Yes
FIIF GroupDro[81]
N/A
R
Yes
N/A CNC[124]
N/A
R
No
N/A GIB[120]
Yes
G
No
FIIF DIR[104]
No
G
No
FIIF CIGA (Ours)
Yes
G
No
PIIF&FIIF
[Figure]
Despite the success of the invariance principle on Euclidean data, the complex nature of graphs raises several new challenges that prohibit direct adoptions of the principle.
First, distribution shifts on graphs are more complicated.
They can happen at both attribute-level and structure-level, and be observed in multiple forms such as graph sizes, subgraph densities and homophily [113, 11, 102].
On the other hand, each of the shifts can spuriously correlate with labels in different modes [4, 71, 2].
Consequently, the entangled complex distribution shifts make it more difficult to identify and capture the invariance on graphs.
Second, OOD algorithms developed and analyzed on Euclidean data often require additional environment (or domain) labels for distinguishing the sources of distribution shifts [4].
However, the environment labels could be highly expensive to obtain and thus often unavailable for graphs, as collecting the labels usually requires expert knowledge due to the abstraction of graphs [37].
These challenges render the problem studied in this paper even more challenging:
How could one generalize the invariance principle to enable OOD generalization on graphs?
To solve the above problem, we propose C ausality
Inspired I nvariant G raph Le A rning (CIGA), a new framework for capturing the invariance of graphs to enable guaranteed OOD generalization under different distribution shifts.
Specifically, we build three Structural Causal Models (SCMs) [74] to characterize the distribution shifts that could happen on graphs: one is to model the graph generation process, and the other two are to model two possible interactions between invariant and spurious features 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 to graphs for OOD generalization: GNN models are invariant to distribution shifts if they focus only on an invariant and critical subgraph G c that contains the most of the information in G about the underlying causes of the label.
Thus, the problem of achieving OOD generalization on graphs can be rephrased into two processes: invariant subgraph identification and label prediction.
Accordingly, shown as Fig. 1(a), we introduce a prototypical invariant graph learning algorithm that decomposes a GNN into: a) a featurizer g for identifying the underlying invariant subgraph G c from G ; b) a classifier f c for making predictions based on G c .
To extract the desired subgraph G c , we derive an information-theoretic objective for the featurizer to identify subgraphs that maximally preserves the invariant intra-class information across a set of different (unknown) environments.
We theoretically show 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 a challenging setting from AI-aided drug discovery [40], show that CIGA can significantly outperform all 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 thorough discussions to Appendix B.2.
On Euclidean data, Invariant Learning [4, 23, 2], Group Distributionally
Robust Optimization [49, 81, 124], Domain Adaption and Domain Generalization [31, 93, 52, 27, 61,
100] are three widely adopted approaches to enable OOD generalization.
However, they all have their own limitations when being applied to graphs.
First, previous invariant learning methods are mostly developed and analyzed for Euclidean data [4, 2, 23], or under specific SCM assumptions [4], making the theoretical results hardly able to generalize to the complicated graph data [80] that can have multiple types of distribution shifts [71].
Group Distributionally Robust Optimization that minimizes the gap between worst group risk and average risk [49, 81, 124], and Domain Adaption/Generalization methods that aim to learn class-conditional domain invariant representations [31, 93, 52, 27, 100], cannot guarantee a min-max optimal predictor without additional assumptions [126, 4, 2].
Moreover, most existing methods require environment labels that are however expensive to obtain in graphs, which limits their applications to graphs [4, 49, 2, 81, 31, 93, 27, 61].
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 input as the explanation for a GNN prediction [116, 122].
Although some may leverage causality to justify the generated explanation [53], they mostly focus on understanding the predictions of GNNs instead of for OOD generalization.
The closest works to ours are two interpretable GNNs that aim to explicitly extract a subgraph for both predictions and explanations guided by information theory [120] and causality [104], respectively.
However, they focus on graphs and shifts generated under a specific
SCM.
Although one of them can provide theoretical guarantee for OOD generalization [120] by using the information bottleneck criteria [2], they would inevitably fail to generalize to graphs generated under different SCMs.
More discussions about the failure are deferred to Appendix D.4.
Besides,
Bevilacqua et al.
[11] also discuss OOD generalization on graphs but limited to a specific graph family and graph size shifts.
Wu et al.
[103] propose OOD generalization algorithms on graphs for the task of node classification, also limited to graphs and shifts under a specific SCM.
To the best of our knowledge, there is no existing work that could handle more comprehensive graph distribution shifts than CIGA, while also achieving provable OOD generalization performance.
OOD Generalization on Graphs through the Lens of Causality 2.1 Problem Setup
In this work, we focus on OOD generalization in graph classification.
Specifically, we are given a set of graph datasets D =
{D e } e collected from multiple environments E all .
Samples ( G ei , Y ei ) ∈
D e from 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 each graph 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 [97, 4].
We leave more details about the background 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 [74, 2].
Thus, we will first formulate the data generation process with structural causal model and latentvariable model [74, 77, 50], to characterize the distribution shifts that could happen on graphs.
Then, we investigate whether the existing methods are generalizable under these distribution shifts.
2.2 Graph Generation Process
[Figure]
We take a latent-variable model perspective on the graph generation process and assume that the graph is generated through a mapping f gen :
Z → G , where Z ⊆ R n is the latent space and G = ∪ ∞ N =
[Equation]
is the graph space.
Let E denote environments.
Following previous works [50, 2], we partition the latent variable from Z into an invariant part 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 or not.
Similarly in images, C and S can represent content and style while E can refer to the locations where the images are taken [7, 125, 50].
Furthermore, C and S control the generation of the observed graphs (Assumption 2.1) and can have multiple types of interactions at the latent space (Assumptions 2.2, 2.3).
Graph generation model.
We elaborate the SCM for the graph generation process in Assumption 2. and Fig. 2(a), where noises in the structural equations are omitted for simplicity [77].
Assumption 2.1 (Graph Generation Structural Causal Model) .
[Equation]
In Assumption 2.1, f gen is decomposed into f G c gen , f G s gen and f G gen to control the generation of G c , G s , and G , respectively.
Among them, G c inherits the invariant information of C that would not be affected by the interventions (or changes) of E [74, 77].
For example, certain properties of a molecule can usually be described by a sub-molecule, or a functional group, which is invariant across different species or assays [12, 92, 40].
On the contrary, the generation of G s and G will be affected by the environment E through S .
Thus, graphs collected from different environments (or domains) can have different distributions of structure-level properties (e.g., graph sizes [11, 102]) as well as feature-level properties (e.g., homophily [62, 17]).
Therefore, the subgraph G s inherits the spurious feature about
Y [125].
In fact, Assumption 2.1 is compatible with many graph generation models by specifying the function classes of f G c gen , f G s gen and f G gen [89, 57, 117, 59].
Since our goal is to characterize the potential distribution shifts in Assumption 2.1, we focus on building a general SCM that is compatible to many graph families and leave graph family specifications and their implications to OOD generalization in future works.
More discussions are provided in Appendix C.
Interactions at latent space.
Following previous works [4, 2], we categorize the latent interactions between C and S into Fully Informative Invariant Features (FIIF, Fig. 2(b)) and Partially Informative
Invariant Features (PIIF, Fig. 2(c)) 2 , depending on whether the latent invariant part C is fully informative about label Y , i.e., ( S, E ) ⊥⊥ Y | C .
Formal definitions of the corresponding SCMs are given as follows, where noises are omitted for simplicity [74, 77].
Assumption
[Equation]
Assumption
[Equation]
In the two SCMs above, f gen corresponds to the graph generation process in Assumption 2.1, and f spu is the mechanism describing how S is affected by C and E at the latent space.
By definition, S is directly controlled by C in FIIF and indirectly controlled by C through Y in PIIF, which can exhibit different behaviors in the observed distribution shifts.
In practice, performances of OOD algorithms can degrade dramatically if one of FIIF or PIIF is excluded [5, 71].
This issue can be more serious in 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, 86, 87], which also serves as the necessary separation assumption for a classification task [69, 16, 65].
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 to generalize to OOD graphs under different distribution shifts, which implies the invariant GNN below 3 .
Definition 2.5 (Invariant GNN) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutation invariant 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 negative unfortunately.
Based on the synthetic BAMotif graph classification task [58, 104] shown in Fig.
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]
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 GNNs trained with the standard empirical risk minimization (ERM) algorithm [97] are not able to generalize to 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 data biases, 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 attributes in the graphs [33].
This phenomenon has also been shown to exist in GNNs equipped with more sophisticated architectures such as attention mechanisms [99], under graph size shifts [46].
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 ) [80].
In addition to IRM, the failure would also happen for alternative objectives [49, 9, 2] as proved by Rosenfeld et al.
Besides, different distribution shifts on graphs can be nested with each other where each one can have distinct spurious correlation type, e.g., FIIF or PIIF.
OOD objectives will also fail seriously if either of the correlation types is not supported [5, 71].
Moreover, non-trivial environment partitions or labels are required for performance guarantee of these OOD objectives [4, 49, 81, 2].
However, collecting meaningful environment partitions of graphs requires expert knowledge about graph data.
Thus, the environment labels can be expensive to obtain and are usually not available [67, 28, 37].
Alternative options such as random partitions tend not to alleviate the issue [23, 55], as it can be trivially deemed as mini-batching.
Challenges of OOD generalization on graphs.
The aforementioned failure analysis reveals that existing methods or objectives fail to elicit an invariant GNN primarily due to the following two challenges:
a) Distribution
shifts on graphs are more complicated where different types of spurious correlations can be entangled via different graph properties; b) Environment labels are usually not available due to the abstraction of graphs.
Despite these challenges, we are still highly motivated to address the following research question: Would it be possible to learn an invariant GNN that is generalizable 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, we generalize 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 set of variables that have stable causal relationship with Y under both FIIF and PIIF (Assumption 2.2, 2.3).
According to the ICM assumption [77], the labeling process C
→ Y is not informed nor influenced by other processes, implying that the conditional distribution P ( Y | C ) remains invariant to the interventions on the environment latent variable E [74].
Consequently, for a GNN with a permutation invariant 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 ρ resembles traditional ERM [97] 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 C involves two causal mechanisms f G c gen and f G gen in Assumption 2.1.
The unavailability of E further adds 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 and f G gen in Assumption 2.1.
The idea of alignment in CIGA is motivated by the algorithmic reasoning results that a neural network can learn a reasoning process better if its computation structure aligns with the process better [108, 110].
Specifically, we realize the alignment by decomposing a GNN into two sub-components 4 : a) a featurizer GNN g : G → G c aiming to identify the desired G c ; b) a classifier GNN f c : G c
→
Y that predicts the label Y based on the estimated G c , where G c refers to the space of subgraphs of G .
Formally, the learning objectives of f c and g can be formulated as:
[Equation]
where maximizing I ( (cid:98) G c ; Y ) is equivalent to minimizing a variational upper bound of R ( f c ( (cid:98) G c ))
[3, 120] that takes (cid:98) G c as inputs to predict label Y for G through f c and g , and (cid:98) G c is the estimated subgraph containing the information about C and hence needs to be independent of E .
Moreover, the extracted G c can either shares the same graph space with input G or has its own space with 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 [122], hence can serve 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 optimizing this architecture cannot satisfy (cid:98) G c ⊥⊥ E .
Formally, merely maximizing I ( (cid:98) G c ; Y ) may include a subgraph from G s in (cid:98) G c since G s also shares certain mutual information with Y .
Moreover, the unavailability of E prevents the direct usage of E in enforcing the independence that is often adopted by previous methods [4, 49, 81, 31, 93], making the identification of G c more challenging.
Optimization objective.
To mitigate this issue, we need to find and translate other properties of G c into some differentiable and equivalent objectives to satisfy the independence constraint (cid:98) G c ⊥⊥ E .
The goal of the desired objective.
We begin by considering a simplistic setting where all the invariant subgraphs G c have the same size s c , i.e., | G c | =
s c 5 .
When maximizing I ( (cid:98) G c ; Y ) in Eq.
1, both FIIF and PIIF can introduce part of G s into (cid:98) G c .
In FIIF (Fig. 2(b)), as G c already contains the maximal possible information in G about Y , G c is a solution to max I ( (cid:98) G c ; Y ) .
However, some subgraph of
G c can be replaced by some subgraph of G s that is equally informative about Y .
In PIIF (Fig. 2(c)), there also exists some subgraph of G s that contains additional information about Y than G c , hence (cid:98) G c is more likely to involve some subgraph of G s .
Thus, the new objective needs to eliminate the auxiliary subgraphs of (cid:98) G c from G s such that the estimated (cid:98) G c can only contain G c .
An important property of G c .
Under both FIIF and PIIF SCMs (Fig. 2), for G e 1 c , G e 2 c that relate to the same causal factor c under two environments e 1 and e 2 , the desired (cid:98) G e 1 c , (cid:98) G e 2 c in e 1 and etend to have high mutual information, i.e., ( G e 1 c , G e 2 c ) ∈ arg max I ( (cid:98) G e 1 c ; (cid:98) G e 2 c ) .
While for G e 1 c and another G e 1 c ′ corresponding to a different c ′ ̸
= c , under the same environment e 1 , including any subgraph from G e 1 s in (cid:98) G e 1 c , (cid:98) G e 1 c ′ will enlarge their mutual information, or in other words, ( G e 1 c , G e 1 c ′ ) ∈ arg min I ( (cid:98) G e 1 c ; (cid:98) G e 1 c ′ ) .
Thus, we can derive an important property of G c , that is, ∀ e 1 , e 2 ∈ E all ,
[Equation]
where (cid:98) G e 1 c
and (cid:98) G e 2 c are the estimated invariant subgraphs corresponding to the same causal factor c under environment e
and e 2 , respectively, while (cid:98) G e 2 c ′ corresponds to a different causal factor c ′ .
Deriving CIGA v1 based on the identified property of G c .
In practice, C is not given.
Nevertheless, since C and Y shares a stable causal relationship in both FIIF and PIIF SCMs, Y can serve as a proxy of C in Eq. 2.
Moreover, as Eq. 2 holds for any ∀ e 1 , e 2 ∈ E all , the environment superscripts can be eliminated without affecting Eq. 2.
Furthermore, when both I ( (cid:98) G e 1 c ; (cid:98) G e 2 c | C = c ) and I ( (cid:98) G c ; Y ) are maximized, I ( (cid:98) G e 1 c ; (cid:98) G e 1 c ′ | C
=
c
′ , c ′ ̸
= c ) is automatically minimized, otherwise all classes will collapse to trivial solutions which is contradictory given I ( (cid:98) G c ; Y ) being maximized.
Therefore, we can derive an alternative objective to Eq.
1 by leveraging Eq.
2 to replace the independence condition: ( CIGAv1 )
[Equation]
where (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G is sampled from training graphs that share the same label Y as G .
In Theorem 3.1, we show how Eq. 3 is equivalent to Eq.
Nevertheless, Eq. 3 requires a strong assumption on 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.
2 without additional constraints will lead to the presence of part of G s in (cid:98) G c .
For instance, (cid:98) G c =
G is a trivial solution to Eq.
3 when s c = ∞ .
Deriving CIGA v2 by resolving size constraint on G c in CIGA v1.
To this end, we further resort to the properties of G s .
In both FIIF and PIIF SCMs (Fig. 2), G s and G c can share certain overlapped information about Y .
When maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) and I ( (cid:98) G c ; Y ) , the appearance of partial G s in (cid:98) G c will not affect the optimality.
However, it can reduce the mutual information between the left part (cid:98) G s =
G − (cid:98) G c and Y , i.e., I ( (cid:98) G s ; Y ) .
Therefore, by maximizing I ( (cid:98) G s ; Y ) , we can reduce including part of G s into (cid:98) G c .
Meanwhile, to avoid trivial solution that G c ⊆ (cid:98) G s during maximizing
I ( (cid:98) G s ; Y ) , we can leverage the better clustering property of G c implied by Assumption 2.4 to derive the constraint I ( (cid:98) G s ; Y ) ≤
I ( (cid:98) G c ; Y ) .
Thus, we can obtain a new objective CIGAv2 as follows:
[Equation]
where (cid:98) G c
=
g ( G ) , (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G is sampled from training graphs that share the same label Y as G .
We also prove the equivalence between Eq.
4 and Eq. 1 in Theorem 3.1.
3.3 Theoretical Analysis and Practical Discussions
Theorem 3.1 (CIGA Induces Invariant GNNs) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c gen in Assumption 2.1 are invertible, (b) samples from each training environment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr , then:
(i).
If ∀ G c , | G c | = s c , then each solution to Eq. 3, elicits an invariant GNN (Def. 2.5).
(ii).
Each solution to Eq. 4, elicits an invariant GNN (Def. 2.5).
We prove Theorem 3.1 (i) and (ii) in Appendix E.2, E.3, respectively.
Practical implementations of CIGA objectives.
After showing the power of CIGA, we introduce the practical implementations of CIGAv1 and CIGAv2 objectives.
Specifically, an exact estimate of the second term I ( (cid:98) G c ; (cid:101) G c | Y ) could be highly expensive [96, 8].
However, contrastive learning with supervised sampling provides a practical solution for the approximation [42, 20, 82, 96, 8]:
[Equation] [Equation]
[Table]
where positive samples ( (cid:98) G c , (cid:101) G c ) are the extracted subgraphs of graphs that share the same label as
G , negative samples are those having different labels, P g ( G |Y =
Y ) is the push-forward distribution of P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label Y , P ( G
|Y ̸ = Y ) refers to the distribution of G given the label that is different from Y , h (cid:98) G c , h (cid:101) G c , h G ic are the graph presentations of the estimated subgraphs, and ϕ is the similarity metric for graph representations.
As M → ∞ , Eq.
5 approximates I ( (cid:98) G c ; (cid:101) G c | Y ) , which can be regarded as a nonparameteric resubstitution entropy estimator via the von Mises-Fisher kernel density [1, 41, 101].
Thus, plugging it into Eq. 3 and Eq.
4 can relieve the issue of approximating I ( (cid:98) G c ; (cid:101) G c | Y ) in practice.
To implement I ( (cid:98) G s ; Y ) given the constraint I ( (cid:98) G s ; Y )
≤
I ( (cid:98) G c ; Y ) in CIGAv2, a practical choice is to adopt hinge loss that implement the constrained I ( (cid:98) G s ; Y ) as 1 N R (cid:98) G s · I ( R (cid:98) G c ≤ R (cid:98) G s ) , where N is the number of samples, I is an indicator function that outputs 1 when the inner condition is satisfied otherwise 0 , and R (cid:98) G s and R (cid:98) G c are the empirical risk vector of the predictions for each sample based on the corresponding (cid:98) G s and (cid:98) G c .
More implementation details can be found in Appendix F.
Discussions and implications of CIGA.
Although using contrastive learning to improve OOD generalization is not new in the literature [27, 61, 124], 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 without the decomposition for causal algorithmic alignment.
However, in the experiments we found that merely using the contrastive objective, i.e., CNC [124], yields unsatisfactory OOD generalization performance, which further implies the necessity of the decomposition in CIGA.
Moreover, the architecture of CIGA can have multiple other implementations for both the featurizer and classifier, such as identifying G c at the latent space [86, 87].
Since we cannot enumerate every possible implementation, in this work we choose interpretable GNN architectures as a prototype validation for CIGA and leave more sophisticated architectures as future works.
In particular, when optimized with ERM objective, CIGA can be reduced to interpretable GNNs .
However, merely using interpretable GNNs such as ASAP [78], GIB [120] or DIR [104] cannot yield satisfactory OOD performance.
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 also empirically 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 [104] where artificial structural shifts and graph size shifts 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 relationships and distribution shifts, we also use DrugOOD [40] from AI-aided Drug Discovery with Assay,
Scaffold, and Size splits, convert the ColoredMNIST from IRM [4] using the algorithm from Knyazev et al.
[46] to inject attribute shifts, and split Graph-SST [122] to inject degree biases.
To compare with previous specialized OOD methods for graph size shifts [113, 11], we use the datasets in Bevilacqua et al.
[11] that are converted from TU benchmarks [67].
More details can be found in Appendix G.1.
Baselines and our methods.
Besides the ERM, we also compare with SOTA interpretable GNNs,
GIB [120], ASAP Pooling [78], and DIR [104], to validate the effectiveness of the optimization
[Table]
objective in CIGA.
We use the same selection ratio (i
[Table]
e., s c ) for all models.
Moreover, to validate the effectiveness of the decomposition in CIGA, we compare CIGA with SOTA OOD objectives including IRM [4], v-Rex [49] and IB-IRM [2], for which we apply random environment partitions following [23].
We also compare CIGA with EIIL [23] and CNC [124] that do not require environment labels, where CNC [124] has a more sophisticated contrastive sampling strategy for combating subpopulation shifts.
More implementation and comparison details are deferred to Appendix G.2.
Evaluation.
We report the classification accuracy for all datasets, except for DrugOOD datasets where we use ROC-AUC following [40], and for TU datasets where we use Matthews correlation coefficient following [11].
We repeat the evaluation multiple times, select models based on the validation performances, and report the mean and standard deviation of the corresponding metric.
For each dataset, we also report the “Oracle” performances that run ERM on the randomly shuffled data.
OOD generalization performance on structure and mixed shifts.
In Table 1, we report the test accuracy of each method, where we omit GIB due to its poor convergence.
Different biases indicate different strengths of the distribution shifts.
Although the training accuracy of most methods converges to more than 99% , the test accuracy decreases dramatically as the bias increases and as more 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 and test 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 lower variance , which demonstrates the effectiveness and excellent OOD generalization ability of CIGA.
[Table]
OOD generalization performance on realistic shifts.
In
Table 2 and Table 3, we examine the effectiveness of CIGA in real-world data and more complicated distribution shifts.
Both averaged accuracy and ranks are reported because of the dataset heterogeneity.
Since the tasks are harder than synthetic ones, interpretable GNNs and OOD objectives perform similar to or even under-perform the ERM baselines, which is also consistent to the observations in non-linear benchmarks [34, 40].
However, both CIGAv1 and CIGAv2 consistently and significantly outperform previous methods, including previous specialized methods Γ
GNNs [11] for combating graph size shifts, demonstrating the generality and superiority of CIGA.
Comparisons with advanced ablation variants.
As discussed in Sec. 3.3, CIGA can be reduced to interpretable GNNs and contrastive learning approaches.
However, across all experiments, we can observe that neither the advanced interpretable GNNs (DIR) nor sophisticated contrastive objectives with specialized sampling strategy (CNC) can yield satisfactory OOD performance, which serves as 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 accurate
[Figure]
t e s t a u c ermCIGAv
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
t e s t a u c ermCIGAv
[Figure]
m
a tt h e w s c o rr e l a t i o n c o e ff i
c i e n t ermCIGAv
[Figure] [Figure]
i c
[Figure]
s c , the improvements in CIGAv1 are not as stable as CIGAv2 or even unsatisfactory when the assumption is violated.
This phenomenon also reveals the superiority of CIGAv2 in practice.
Hyperparameter sensitivity analysis.
To examine how sensitive CIGA is to the hyperparamters α and β for contrastive loss and hinge loss, respectively.
We conduct experiments based on the hardest datasets from each table (i.e., SPMotif-Mixed with the bias of 0 .
9 , DrugOOD-Scaffold and the NCI109 datasets from Table 1, Table 2, and Table 3, respectively.)
with different α and β .
When changing the value of β , we fix the α to a specific value under which the model has a relatively good performance (but not the best, to fully examine the robustness of CIGA in practice).
The results are shown in Fig.
4 and Fig.
It can be found that both CIGAv1 and CIGAv2 are robust to different values of α and β , respectively, across different datasets and distribution shifts.
Besides, the results also reflect the effects of the additional penalty terms in CIGA.
For example, in Fig.
16, when α is too small, the invariance of the identified invariant subgraphs (cid:98) G c may not be guaranteed, resulting worse performances.
Similarly, as shown in Fig. 17, when β becomes too small, some part of the spurious subgraph may still appear in the estimated invariant subgraphs, which yields worse performances.
Besides, when α and β become too large, the optimization of CIGA can be affected due to their intrinsic conflicts with ERM, hence a better optimization scheme for CIGA can be a promising future direction [18].
We provide more details and additional analysis on the efficiency of
CIGA and single environment
OOD generalization performance of CIGA in Appendix G.4, as well as the visualization examples of the identified invariant subgraph in Appendix G.5.
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 promising theoretical and empirical OOD generalization ability under a variety of distribution shifts.
Acknowledgments and Disclosure of Funding
We thank the reviewers for their valuable comments.
This work was supported by GRFfrom the RGC of HKSAR and CUHK direct grant 4055146.
TL was partially supported by Australian
Research Council Projects DP180103424, DE-190101473, IC-190100031, DP-220102121, and
FT-220100318.
YZ and BH were supported by the RGC Early Career Scheme No. 22200720, NSFC
Young Scientists Fund No. 62006202, Guangdong Basic and Applied Basic Research Foundation No.
2022A1515011652, and Tencent AI Lab Rhino-Bird Gift Fund.
E Theory and Discussions
E.1 More discussions on Definition 2.5 for Invariant GNNs . .
E.2 Proof for theorem 3.1 (i) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.3 Proof for theorem 3.1 (ii) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Details of Prototypical CIGA Implementation
G Detailed Experimental Settings
G.1 Details about the datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.2 Training and Optimization in Experiments . . . . . . . . . . . . . . . . . . . . . .
G.3 Software and Hardware . . . .
G.4 Additional Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.5 Interpretation Visualization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A Broader Impacts
Considering the wide applications and high sensitivity of GNNs to distribution shifts and spurious correlations, it is important to develop GNNs that are able to generalize to OOD data, especially for realistic 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 towards tackling 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 not involve any human subjects, practices to data set releases, potentially harmful insights, methodologies and 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 Work and Future Directions B.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 feature dimension 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 that learns a meaningful representation h G for each graph G to help predict their labels y G
= ρ ( h G ) with a downstream classifier ρ : R h
→ Y .
The representation h G is typically obtained by performing pooling with a READOUT function on the learned node representations:
[Equation]
where the READOUT is a permutation invariant function (e.g., SUM, MEAN) [107, 115, 70, 107, 19, 68], and h ( K ) u stands for the node representation of u ∈ V at K -th layer
that is obtained by neighbor aggregation:
[Equation]
where N ( u ) is the set of neighbors of node u , σ ( · ) is an activation function, e.g., ReLU, and a ( · ) is an aggregation function over neighbors, e.g., MEAN.
Invariant Learning.
Invariant learning typically considers a supervised learning setting based on the data D
= {D e } e collected from multiple environments E all , where D e = { G ei , y ei } is the dataset from environment e ∈ E all .
( G ei , y ei ) from a single environment e are considered as drawn independently from an identical distribution P e .
The goal of OOD generalization is to train a GNN ρ ◦ h : G → Y with 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 [97, 76, 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 graph as the explanation for the prediction of a GNN model [116, 122].
Although some may leverage causality in explanation generation [53], they mostly focus on understanding the predictions of
GNNs in a post-hoc manner instead of OOD generalization.
Recently there are two works aiming to provide robust explanations under distribution shifts, i.e., GIB [120] and DIR [104], and both of them focus on tackling FIIF spurious correlations (Assumption C.2).
The theoretical guarantees of GIB follows the theory of information bottleneck [95], 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 of them could result in severe performance degradation in practice [4, 2, 5, 71].
While for DIR, though as 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 [84, 85], physics [6, 83], and graph algorithms [94, 98, 108, 105].
Xu et al.
[110] study the neural network extrapolation ability from a geometrical perspective.
Han et al.
[36] improve OOD drug discovery by mitigating the overconfident misprediction issue.
Knyazev et al.
[46], Yehudai et al.
[113] focus on the extrapolation of GNNs in terms of graph sizes, while making additional assumptions on the knowledge about ground truth attentions and access to test inputs.
Bevilacqua et al.
[11] study the graph size extrapolation problem of GNNs through a causal lens, while the induced invariance principle is built upon assumptions on the specific family of graphs.
Different from these works, we consider the GNN extrapolation as a causal problem, establish generic SCMs that are compatible with several graph generation models, as well as, more importantly, different types of distribution shifts.
Hence, the induced the invariance principle and provable algorithms built upon the SCMs in our work can generalize to multiple graph families and distribution shifts.
Additionally, Wu et al.
[103] propose causal models as well as specialized objectives to extrapolate nodes with different neighbors.
However, their formulation is limited to node classification task and specific spurious correlation type.
In contrast, the induced invariance principle in Wu et al.
[103], can be seen as a extension of CIGA for node classification, where we cab identify an invariant subgraph 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 improving the huge success of machine learning algorithms to further advances [75, 86, 87].
One of the most widely applied concept from causality is the Independent Causal Mechanism (ICM) that assumes conditional distribution of each variable given its causes (i.e., its mechanism) does not inform or influence the other conditional distributions [74, 77].
The invariance principle is also induced from the 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 principle to machine learning models for finding an invariant representation about the causal relationship between the underlying causes and the label [76, 4].
Consequently, models built upon the invariant representation can generalize to unseen environments or domains with guaranteed performance [76, 79, 4, 81, 10, 48, 34, 49, 23, 2].
The arguably first formulation of invariance principle was introduced by Peters et al.
Arjovsky et al.
[4] propose a novel formulation of learning causal invariance in representation learning, i.e., IRM, show how it connects with existing areas such as distributional robust optimization [72] and generalization [123], and prove its effectiveness in addressing PIIF spurious 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, 71].
Ahuja et al.
[2] then propose to add information bottleneck criteria into the IRM formulation to address the issue.
However, their results are restricted to linear regime and also require environment partitions to distinguish the sources of distribution shifts.
Recently, Creager et al.
[23] and Lin et al.
[55] propose new OOD objectives to relieve the needs for environment partitions, but limited to PIIF spurious types and linear regime.
Besides, Lin et al.
[54] identify the overfitting problem as a key challenge when applying IRM on large neural networks.
Zhou et al.
[127] propose to alleviate this problem by imposing sparsity constrain.
In parallel invariant learning approaches, Sagawa* et al.
[81] propose to regularize the worst group in group distributionally robust optimization (GroupDro).
Zhang et al.
[124] propose a contrastive approach to tackle GroupDro when the group partitions are not available.
However, minimizing the gap between worst group risk and averaged risk can not yield a OOD generalizable predictors in our circumstances.
Besides, traditional approaches to tackle OOD generalization also include Domain
Adaption, Transfer Learning and Domain Generalization[79, 21, 31, 93, 52, 27, 61, 100], which aim to 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 the ground truth predictors [34, 2], hence are not able to yield predictors with OOD generalization guarantees.
We refer interested readers to Pearl [75], Sch¨olkopf [86], Sch¨olkopf et al.
[87] for an in-depth understanding, and Gulrajani and Lopez-Paz [34], Ahuja et al.
[2] for a thorough overview.
B.3 More discussions on connections of CIGA with existing work
Although primarily serving for graph OOD generalization problem, our theory complements the identifiability study on graphs through contrastive learning, and aligns with the discoveries in the image domain that contrastive learning learns to isolate the content ( C ) and style ( S ) [128, 50].
Moreover, our results also partially explain the success of graph contrastive learning [118, 60, 119], 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 constrained by the encoders embedded for learning graph representations.
During isolating G c from G , if the encoder can not differentiate two isomorphic graphs G c and G c ∪ G ps where G ps ⊆ G s , then the featurizer will fail to identify the underlying invariant subgraph.
Moreover, the classifier will also fail if the encoder can not differentiate two non-isomorphic G c s from different classes.
Thus, adopting more powerful graph representation encoders into CIGA can improve the OOD generalization.
On CIGA and graph information bottleneck.
Under the FIIF assumption on latent interaction, the independence condition derived from causal model can also be rewritten as Y ⊥⊥ S | C (similar to that in DIR [104] as they also focus on FIIF), which further implies Y ⊥⊥ S | (cid:98) G c .
Hence it is natural to use
Information Bottleneck (IB) objective [95] to solve for G c :
[Equation]
which explains the success of many existing works in finding predictive subgraph through IB [120].
However, the estimation of I ( (cid:98) G c , G ) is notoriously difficult due to the complexity of graph, which can lead to unstable convergence as observed in our experiments.
In contrast, optimization with contrastive 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 graph classification based on the ego-graphs of a node, our analysis and discoveries can generalize to node classification.
More specifically, the invariance principle for node classification can be implemented by identifying an invariant subgraph 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 [103].
B.4 Discussions on limitations of CIGA and future directions
Better graph generation modeling.
Compared to Bevilacqua et al.
[11], we do not specify a specific graph family in the SCM for graph generation process.
Since our focus is to describe the potential distribution shifts with SCMs, in Assumption 2.1, we aim to build a SCM that is compatible to many graph generation processes [89, 57, 117, 59].
However, it is often the case that practitioners have certain inductive knowledge about the graph generation process, which may imply useful leads and invariance in modeling the generation process [111, 30, 56].
In Appendix C.1, we provide an example about incorporating the graphon [57] knowledge into the SCMs, which derives similar solutions as in the literature [113, 11].
Therefore, we believe it is promising to leverage more additional knowledge for more precise graph generation modeling and better OOD generalization on graphs.
Better contrastive sampling.
Typical contrastive or graph contrastive learning approaches leverage augmentation techniques as well as sophisticated sampling strategies during the positive or negative pairs selection [20, 82, 96, 118, 119].
A better augmentation or sampling strategy can benefit the
OOD generalization in general as shown by K¨ugelgen et al.
[50] and Zhang et al.
Since our implementation of CIGA in this work aims to verify the theoretical findings, we do not apply sophisticated augmentation or sampling during the sampling while simply using the supervised contrastive approach [42].
Nevertheless, it is promising to leverage better augmentation and contrastive strategy to improve the generalization ability in CIGA [121].
More sophisticated architectures/parameter tunning.
The CIGA framework introduced in Sec.can have multiple implementations.
We choose interpretable architectures in our experiments for the purpose of concept verification.
Essentially, different architectures can have different advantages and limitations.
For the interpretable GNNs used in our experiments, it can provide interpretability for the results (as shown in Appendix G.5), but still requires more training time (as shown in Appendix G.4).
Therefore, it may not be applicable to some resource-limited scenarios such as Edge-AI.
Besides, the approximation may also be limited to the chosen architectures.
More sophisticated architectures can be incorporated, such as identifying and disentangling G c at the latent space [86, 87].
Moreover, as shown in Appendix G.4, CIGA still requires certain additional tunning efforts for the objectives.
Hence we believe it is also a promising future direction to reduce the parameter tunning by leveraging better optimization techiniques [88, 18]
C Full Structural Causal Models on Graph Generation
Due to the space constraints in the main paper, we make some simplifications when giving the SCMs on the graph generation process.
Hence in this section, supplementary to the graph generation process in Sec.
2.2, we provide full SCMs on the graph generation process in this section as shown in Fig. 6.
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 and assume that the graph is generated through a mapping f gen : Z → G ,
where Z ⊆ R n is the latent space and G = ∪ ∞ N =1 { 0 , 1 } N × R N × d is the graph space.
Let E denote environments.
Following previous works [50, 2], we partition the latent variable from Z into an invariant part 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 are taken [7, 125, 50].
While in graphs, C can be the latent variable that controls the generation of functional groups in a molecule, which can not be affected by the changes of environments, such as species (or scaffolds), experimental environment for examining the chemical property (or assays) [40].
On the contrary, the other latent variable S inherits environment-specific information thus can further affect the finally generated graphs.
Besides, C and S can have multiple types of interactions at the latent space with environments E and labels Y , which will generate different types of spurious correlations [2].
Assumption C.1 (Graph generation SCM) .
[Equation]
Specifically, the graph generation process is shown as Fig. 6(a).
The generation mapping f gen is decomposed 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 , they control the generation of the adjacency matrix and features for the invariant subgraph G c and spurious subgraph 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, and subgraph densities.
While Z c X and Z s X mainly control the attribute-level properties in the generated graphs, 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 processes controlled by the latent variables [89, 57, 117, 59, 11].
Note that since our focus is to describe the potential distribution shifts with SCMs, in Assumption 2.1, we aim to build a SCM that is compatible to many graph generation processes [89, 57, 117, 59].
In fact, in Appendix C.1, we showcase how our
SCMs can generalize to specific graph families studied in the literature [11, 104, 103], when given more additional knowledge about the graph generation process.
Nevertheless, we believe integrating specific graph generation processes and their implications to improving OOD generalization on graphs would be a promising future direction, as discussed in Appendix B.4.
Due to the correlation between E and G , graphs collected from different environments can have different structure-level properties such as degrees, graph sizes, and subgraph densities, as well as feature-level properties such as homophily [46, 113, 11, 17].
Meanwhile, all of them can spuriously correlated with the labels depending on how the underlying latent variables are interacted with each others.
The interaction types can be further divided into two axiom types FIIF and PIIF, as well as the mixed one MIIF.
Previous OOD methods such as GIB [120] and DIR [104] mainly focus on FIIF case, while others such as IRM [4] mainly focuses on the PIIF case.
Evidences show that
S C
Z cX
Z cA
Z sX
Z sA
G c
[Figure]
s
[Figure]
E
S Y
G
C
[Figure]
SCM
E
S Y
G
C
[Figure]
SCM
E
S
C
S 2 Y
G
[Figure]
MIIF SCM failing to model either of them when developing the OOD objectives can have serious performance degenerations in practice [5, 71].
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. 6(b)), and Partially Informative Invariant Features (PIIF, Fig. 6(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. 6(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 [74, 77].
Since MIIF is built upon FIIF and PIIF, we will focus on the axiom interaction modes (FIIF and PIIF) in this paper, while most of our discussions can be extended to MIIF or more complex interactions built upon FIIF and PIIF.
Among all of the interaction modes, f gen corresponds to the graph generation process in Assumption C.1.
f spu is the mechanism describing how S is affected by C and E at the latent space.
In
FIIF, S is directly controlled by C while in PIIF, indirectly controlled by C through Y , which can exhibit different behaviors in practice [2, 71].
Additionally, in MIIF, S is further partitioned into Sand 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 merely based on C .
Consequently, C is better clustered than S when given Y [13, 15, 86, 87], which also serves as the necessary separation assumption for a classification task [69, 16, 65].
Assumption C.5 (Latent Separability) .
H ( C | Y ) ≤ H ( S | Y ) .
C.1 Discussions on specific cases of the SCMs
Although our primary focus in this work is to characterize general graph distribution shifts that could happen in practice without any additional knowledge about the underlying graph family, and derive the corresponding solutions, our SCMs (Fig. 6) can generalize to specific cases studied in previous works, when incorporating more inductive biases about the underlying graph family [11, 104, 103].
Specifically, we illustrate the specialized SCMs in Fig.
7 for the SCM studied in [11] which assumes the graphs are generated following a graphon model [57].
When with the additional knowledge about the underlying graph generative model, the graph generation SCM (Fig. 6(a)) and the FIIF SCM (Fig. 6(b)) together generalizes to the graphon SCM studied in [11].
We now give a brief description in the below.
[Figure]
GraphonModel
[Figure]
Graphon SCM from [11].
Specifically, shown as in Fig. 7(a), C now is instantiated as a graphon model C W ∼ P ( C W ) , where
C W : [0 , 1] 2 → [0 , 1] is a random symmetric measurable function sampled from the set of all symmetric measurable functions [57].
Besides, the label Y is determined according to C W .
Then,
C W will further control the generation of the adjcency matrix G c
=
A c through graphon generative process:
[Equation]
where Z u,v is an independent uniform noises on [0 , 1] for each possible edge ( u, v ) in the graph.
Bascially, Z and U are inherited from the graphon SCM as Fig. 7(c).
On the other hand, as S does not imply any information about Y in this case, it resembles the FIIF
SCM (Fig. 6(b)).
In other words, ( S, E ) ⊥⊥ Y | C still holds.
Moreover, the node attributes G s =
X s are generated jointly influenced by the environment E and the graphon C W through S :
[Equation]
which resembles the attribute generation in Fig. 7(c).
Then, both G c and G s are concatenated together.
In a simplistic case intuitively, we can regard G c only contains the edges in G and G s only contains the node attributes.
Since the graphon model mainly controls the edge connection, the edge connection patterns, e.g., motif appearance frequency or subgraph densities, acts as a informative indicator for the label Y .
In contrast, the node attributes and its numbers would be affected by the environments.
A GNN model is prone to the changes of the environments if it overfits to some spurious patterns about the graph sizes or the attributes.
While if the GNN model can leverage the connection patterns to make predictions, it remain invariant to the changes of environments, or the spurious patterns such as graph sizes and node attributes, which resembles the solutions derived in [113, 11].
Besides, it also partially explains why CIGA can generalize to OOD graphs studied in these works [113, 11].
In addition to the graphon SCM, essentially, the SCM studied in [104] resembles the FIIF SCM, and that of [103] resembles PIIF SCM, which also serves as partial evidence for the superiority OOD generalization performances of CIGA.
D More Details about Failure Case Studies in Sec. 2.
In this section, we provide details on failure case studies in Sec.
We first elaborate the empirical evaluation setting where we construct a synthetic graph datasets to probe the behaviors of existing methods in OOD generalization on graphs.
D.1 More empirical details about failure case study in Sec. 2.
To begin with, we construct 3-class synthetic datasets based on BAMotif [58] and follow Wu et al.
[104] to inject spurious correlations between motif graph and base graph during the generation.
In this graph classification task, the model needs to tell which motif the graph contains, e.g., “House” or “Cycle” motif, as shown in Fig.
We inject the distribution shifts in the training data while keeping the test data and validation data without the biases.
For structure-level shifts, we introduce the artificial bias based on FIIF, where the motif and the base graph are spuriously correlated with a probability of various bias.
For mixed shifts, we additionally introduced attribute-level shifts based on FIIF, where all of the node features are spuriously correlated with a probability of various bias.
The number of training graphs is 600 for each class and the number of graphs in validation and test set is 200 for each class.
More construction details are given in Appendix G.
For the GNN encoders, by default, we use 3 -layer GCN [45] with mean readout, a hidden dimension of 64 , and JK jump connections [106] at the last layer.
During training, we use a batch size of 32 , learning rate of 1 e − 3 with Adam optimizer [43], and batch normalization between hidden layers [39].
Meanwhile, to stabilize the training, we also use dropout [91] of 0 .
1 and early 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 reported and plotted.
When using
IRM objective [4], as the environment partitions are not available, we generate 2 environments with random partitions.
D.2 More discussions about failure case study in Sec. 2.
[Figure]
ERM can fail for leveraging the shortcuts and predict graphs that have a hexagon or have mostly green nodes as “house”.
IRM can fail when test data is not sufficiently supported by the training data.
In Fig.
9, 10, 11, 12, we investigate whether existing training objectives (ERM and IRM), adding more message passing, as well as using expressive GNNs, can improve the OOD generalization ability on graphs.
Here we also provide a additional discussion in complementary to the discussions on OOD generalization 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 signals by aggregating more neighbors with higher layers, or enhance the expressive power with more powerful readout functions [106, 107, 112].
Aggregating neighbor information with more layers to denoise the input signal, or enhancing the expressivity with more powerful readout functions, are two common choices in GNNs to improve the generalization ability [106, 51, 107, 112].
However, in the experiments next, we empirically found that GCNs with more layers and more powerful readout operations are still sensitive to distribution shifts.
In particular, stacking more layers helps denoising certain shifts, while the OOD performance would drop more sharply when the bias increases.
Intuitively, if the spurious features from nodes cannot be eliminated by the denoising property of a deeper GNN, they would spread among the whole graph more widely, which in turn leads to stronger spurious correlations.
Besides, the spurious correlations would be more difficult to be disentangled if there are distribution shifts at both structure-level and attribute-level.
Since the node representations from hidden layers can also encode graph topology features [107], distribution shifts introduced through Z sA and Z sX will doubly mix at the learned features.
In the worst case, the information about Z cA and Z cX could be partially covered by or even replaced by Z sA and Z sX .
This will make OOD generalization of message passing GNNs trained through ERM much more difficult or even impossible.
Besides, as the node representations of 1 ≤ i ≤ k -th layer can also encode graph topology features [107], which, if spuriously correlated with labels through Z sA and entangled with part of invariant node features, i.e., Z cX , in the worst case, can greatly improve the difficulty or even make the OOD generalization impossible for neighbor aggregation GNNs trained with ERM.
Using more expressive GNNs.
Previous results on the expressivity of GNNs show that GNNs are limited to distinguish isomorphic graphs at most as 1-WL/2-WL test can distinguish [107].
After that, many follow-up variants are proposed to improve the expressivity of GNNs [68].
However, if the labels are spuriously correlated with certain subgraphs, even the GNN has high expressivity can still be prone to distribution shifts.
In a idealistic case, when classifying a graph with a highly expressive GNN, it reduces to the linear or discrete feature case on the Euclidean regime.
In this case
[Figure]
there exists many evidences showing that neural networks can fail to generalize to OOD data without a proper objective [7, 24, 4, 81, 10, 49, 23, 48, 2].
Empirically, we use k -GNNs [66] to verify the intuition and observe similar failures for this provably more expressive GNN as basic GNN variants
[Figure]
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.
[104] to construct dataset for Fig.
9, 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.
10, graph size shifts are added, which is exactly the SPMotif datasets used in
DIR [104].
Besides, in Fig.
11, we can also change the FIIF attribute-level shifts to PIIF attributelevel shifts, where we flip the labels by a probability of 5% and let the flipped label to be spuriously correlated with the node features, following the PIIF SCM in Fig.
Graph size shifts can also be injected 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.
9, 10, 11, 12, we can observe that all GNNs are sensitive to distribution shifts.
As the intensity of spurious correlation grows, GNNs are more likely to overfit to shortcuts presented either in the structure-level or attribute-level, which is similar to general deep learning models [33].
Observation II: Higher variance also indicates unstable OOD performance.
Although GNNs show certain robustness against single distribution shifts, e.g., performances do not decrease sharply at the beginning in Fig.
9, when the spurious correlation grows stronger, the OOD performance become 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 simple compared to reality.
Hence the performance will be highly sensitive to the quality of initialized points at 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 distribution shifts, 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.
6, 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.
9, we can find that, when the attribute-level distribution shifts are mixed, the performance will be worse and more unstable.
When the graph size shifts are mixed, this phenomenon will be more obvious, as shown in Fig.
This phenomenon also verifies the observations in Knyazev et al.
[46] that attention mechanism in GNN is also sensitive to graph size shifts and can hardly learn the desired attention distributions without further guidance.
Moreover, when the structure-level and attribute-level shifts have different spurious correlation types, i.e., when FIIF structure-level shifts and PIIF attribute-level shifts are both presented, the performance drop will be more serious, by comparing Fig.
9 to Fig.
11, as well as Fig.
10 to Fig.
Observation IV: Using more powerful architectures can not improve the OOD performance.
From the sub-figures (b) and (c) in Fig.
9, 10, 11, 12, we can also observe that neither adding more message passing turns nor using more expressive GNN architectures can be immune to distribution shifts.
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 they are more likely to learn the desired information during the optimization [109].
However, when the intensity of spurious correlation grows stronger, deeper GNNs are more likely to overfit to shortcuts hence 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.
9, 10, 11, 12 we use 1 - 2 - 3GNN following the algorithm of k -GNNs which is provably more expressive than 2 -WL test [66].
When there are no graph size shifts, k -GNNs will have higher performance at the beginning.
When there are graph size shifts, k -GNNs will have a lower initial performance at the beginning.
Then, as the spurious strength grows, k -GNNs can suddenly become seriously unstable, though k -GNNs can have higher averaged performance, which reflects unsatisfactory OOD performance as Observation II implies.
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 architectures to help improve the OOD generalization performance.
Beyond the empirical studies in previous section, we aim to accompany more formal discussions for explaining the failures of existing optimization objectives and architectures in the next sections.
D.4 Theoretical discussions for failure case study in
Sec.
A motivating example.
To begin with, we follow Ahuja et al.
[2] to introduce a formal example on the failures of GNNs optimized with ERM or IRM [97, 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 the corresponding invariant and varying latent variables, N is Bernoulli binary noise with a parameter of q 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 mean as 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 Definition 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 as above.
When ∪ e ∈E te supp ( P e ) ̸⊆ ∪ e ∈E tr supp ( P e ) , it implies the S from training environments E tr does not cover S from testing environments, while C can be covered.
Moreover, the condition of strictly separable 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 training data, 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 are strictly 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 generalization problem.
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 solution that 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, assuming the graph label is determined by the node label and all of the nodes have the same label that are determined as Definition D.1, then GNN optimized with ERM and IRM will also fail because of the same reasons as case (i).
Discussions on the failures of previous OOD related solutions.
First of all, for IRM or similar objectives [81, 49, 2, 9] that require environment information or non-trivial data partitions, they can hardly be applied to graphs due to the lack of such information.
The reason is that obtaining such information can be expensive due to the abstraction of graphs.
Moreover, as proved in Theorem 5. of Rosenfeld et al.
[80], when there is not sufficient support overlap between training environments and testing environments, the IRM or similar objectives can fail catastrophically when being applied to non-linear regime.
The only OOD objective EIIL [23] that does not require environment labels, also rely on similar assumptions on the support overlap.
We also empirically verify their failing behaviors in our experiments.
Moreover, since part of explainability works also try to find a subset of the inputs for interpretable prediction robustly against distribution shifts.
Here we also provide a discussion for these works.
The first work following this line is I NV RAT [14], which develops an information-theoretic objective (we re-formulate it to suit with OOD generalization problem on graphs):
[Equation] [Equation]
However, it also requires extra environment labels for optimization that are often unavailable in graphs.
Besides, the corresponding assumption on the data generation for guaranteed performance is essentially PIIF if applied to our case, while it can not provide any theoretical guarantee on FIIF.
We also notice a recent work, DIR [104], as a generalization of I NV RAT to graphs while studying
FIIF spurious correlations, that proposes an alternative objective which does not require environment label:
[Equation]
However, the theoretical justification established for DIR (Theorem 1 to Corollary 1 in Wu et al.
[104]) 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 FIIF nor PIIF.
In experiments, we empirically find the unstable and relatively high sensitivity of DIR to spurious correlations, which verifies our finding.
More details about empirical behaviors of DIR can be found in Appendix G.
In contrast to DIR, GIB [120] 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 from causal model can also be rewritten as Y ⊥⊥ S | C (similar to that in DIR [104] as they also focus on FIIF), which further implies Y ⊥⊥ S | (cid:98) G c .
Hence it is natural to use Information Bottleneck (IB) objective [95] to solve for G c :
[Equation]
which explains the success of many existing works in finding predictive subgraph through IB [120].
However, the estimation of I ( (cid:98) G c , G ) is notoriously difficult due to the complexity of graph, which can lead to unstable convergence as observed in our experiments.
In contrast, optimization with contrastive objective in CIGA as Eq.
5 induces more stable convergence.
D.5 Challenges of OOD generalization on graphs.
From the aforementioned analysis, we can summarize some key challenges revealed by the failures of both existing optimization objectives and GNN architectures.
In particular, we are facing two main challenges a) Distribution shifts on graphs are more complicated where different types of spurious correlations can be entangled via different graph properties; b) Environment labels are usually not available due to the abstract graph data structure.
E Theory and Discussions
In this section, we provide proofs for propositions and theorems mentioned in the main paper.
E.1 More discussions on Definition 2.5 for Invariant GNNs
Definition 2.5 is motivated by applying the invariance principle to the established SCMs in Sec. 2.2, following the literature of invariant learning [76].
In this section, we will present Proposition E.2 and
Proposition E.3 to illustrate how satisfying the minmax objective in Definition E.1 is equivalent to identifying the underlying invariant subgraph G c that contains all of the information about causal factor C in G , under both FIIF and PIIF SCMs (Fig. 2(b) and Fig. 2(c)).
Definition E.1 (Invariant GNN) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutation invariant 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 .
First, we show that using the invariant subgraphs G c to predict Y can satisfy the minmax objective min max e ∈E all R e in Proposition E.2.
Proposition E.2.
Let G c denote the subgraph space for G c , given a set of graphs with their labels
D =
{ G ( i ) , y ( i ) }
Ni =1 and E all that follow the graph generation process in Sec.
2.2 (or Sec.
C), a GNN ρ ◦ h : G c → Y that takes G c of G as the input to predict Y , and solves the following objective can generalize to OOD graphs, i.e., solving the minmax objective in Def.
E.1:
[Equation]
where R G c is the empirical risk over { G ( i ) c , y ( i ) } Ni =1 and G ( i ) c is the underlying invariant subgraph
G c for G ( i ) .
Proof.
We establish the proof with independent causal mechanism (ICM) assumption in SCM [74, 77].
In particular, given the data generation assumption, i.e., for both FIIF (Assumption 2.2) and PIIF (Assumption 2.3), we have: ∀ e,
[Equation]
where we use ICM for the first three equalities.
From Eq.
14, it suffices to know P ( Y | G c ) is invariant across different environments.
Hence, a GNN predictor ρ ◦ h : G c → Y optimized with empirical risk given G c , essentially minimizes the empirical risk across all environments, i.e., min R G c = min max R e .
Thus, if ρ ◦ h solves min R G c , it also solves min max R e , hence it elicits a invariant GNN predictor according to Definition.
E.1.
Besides, we show in Proposition E.3 that only using the underlying invariant subgraphs G c to make predictions can satisfy the minmax objectives.
Or equivalently, a GNN predictor solving the minmax objective can only rely on the underlying invariant subgraph G c to predict Y .
Proposition E.3.
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutation invariant graph encoder h : G → R h and a downstream classifier ρ : R h
→ Y , ρ ◦ h that minimizes the worst case risk among all environments, i.e., min max e ∈E all R e , can not rely on any part of G s , i.e., ρ ◦ h ( G ) ⊥⊥ G s .
Proof.
The proof for Proposition E.3 is straightforward.
Assuming that ρ ◦ h ( G ) ̸⊥⊥ G s , as E is influenced by the changes of E through S in both FIIF and PIIF SCMs (Fig. 2(b) and Fig. 2(c))
, then ρ ◦ h ( G ) ̸⊥⊥ E as well.
Consequently, there exists some graph G corresponding to G c , G es and ρ ◦ h ( G ) =
Y under an environment e , such that we can always find a proper e ′ to make ρ ◦ h ( G ) ̸
= Y .
In contrast, the prediction of a GNN that satisfies
ρ ◦ h ( G ) ⊥⊥ G s remains invariant against arbitrary changes of environments.
Thus, it leads to a contradiction to the condition that min max e ′ ∈E all R e ′ .
Therefore, a GNN that solves min max e ∈E all R e must satisfy ρ ◦ h ( G ) ⊥⊥ G s .
Combining Proposition E.2 and Proposition E.3, we are highly motivated to find the underlying invariant subgraphs to make predictions about the original graphs, which converges to Eq. 1.
Tackling
Eq.
1 under the unavailability of E brings us two variants of CIGA solutions, as illustrated in
Section 3.
E.2 Proof for theorem 3.1 (i)
Theorem E.4 (CIGAv1 Induces Invariant GNNs) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c gen in Assumption 2.1 are invertible, (b) samples from each training environment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr , if ∀ G c , | G c | =
s c , then a GNN f c ◦ g solves Eq.
4, is an invariant GNN (Def. 2.5).
Proof.
We re-write the objective as follows:
[Equation]
where (cid:98) G c
=
g ( G ) , (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G and G have the same label.
The proof of Theorem E.4 is essentially to show the estimated (cid:98) G c through Eq.
15 is the underlying
G c , then the maximizer of I ( (cid:98) G c ; Y ) in Eq.
15 can produce most informative and stable predictions about Y based on G , hence is an invariant GNN (Definition.
E.1).
In the next, we are going to take an information-theoretic view of the first term
I ( (cid:98) G c ; Y ) and the second term I ( (cid:98) G c ; (cid:101) G c | Y ) to conclude the proof.
We begin by introducing the following lemma:
Lemma E.5.
Given the same conditions as Thm.
E.4, I ( (cid:98) G c ; Y ) is maximized if and only if
I ( (cid:98) G c ; Y | E = e ) is maximized, ∀ e ∈ E tr .
The proof for Lemma E.5 is straightforward, given the condition that samples from each training environment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr .
Obviously, (cid:98) G c =
G c is a maximizer of I ( (cid:98) G c ; Y )
=
I ( C ; Y ) =
H ( Y ) , since f c gen : C → G c is invertible and C causes
Y .
However, there might be some subset G p s ⊆
G s from the underlying G s that entail the same information about label, i.e., I ( G pc ∪ G ps ; Y ) =
I ( G c ; Y ) where (cid:98) G c
=
G pc ∪ G ps and G pc = G c ∩ (cid:98) G c .
For FIIF (Assumption 6(b)), it can not happen, otherwise, let G lc
=
G c − G pc , then we have:
[Equation]
where the second last equality is due to C → Y and the invertibility of f c gen : C → G c in FIIF, i.e.,
H ( Y | C ) =
H ( Y | G c )
=
H ( Y | G pc , G lc ) = 0 .
However, in PIIF, it can hold since conditioning on
G pc , G ps can not determine Y , as S ̸⊥⊥ Y | C .
In other words, G s ̸⊥⊥ Y | G c , which means G s can imply some information about Y that is equivalent to I ( G lc ; Y | G pc ) .
To avoid the presence of spuriously correlated G s in (cid:98) G c , we will use the second term to eliminate it:
[Equation]
where (cid:98) G c
=
g ( G ) , (cid:101) G c
=
g ( (cid:101) G ) are two positive samples drawn from the same class (i.e., condition on the same Y ).
Since the all of the training environments are equally distributed, maximizing
I ( (cid:98) G c ; (cid:101) G c | Y ) is essentially maximizing I ( (cid:98) G c , E = ˆ e ;
[Equation]
G c , E = ˜
e | Y ) , ∀ ˆ e, ˜ e ∈ E tr .
Hence, we have:
[Equation]
We claim Eq.
18 can eliminate any potential subsets from G s in the estimated (cid:98) G c .
[Equation]
(cid:98) G ls , respectively.
Otherwise, suppose there are some subsets (cid:98) G p s ⊆ (cid:98) G s and (cid:101)
G p s ⊆ (cid:101)
G s contained in the estimated (cid:98) G c , (cid:101) G c , where (cid:98) G s , (cid:101) G s be the corresponding underlying G s s for (cid:98) G c , (cid:101) G c .
Let (cid:98) G ∗ c and (cid:101) G ∗ c be the ground truth invariant subgraph G c s of (cid:98) G and (cid:101) G , (cid:98) G lc
=
(cid:98) G ∗ c − (cid:98) G c
and (cid:101) G lc
=
(cid:101) G ∗ c − (cid:101) G c be the l eft (un-estimated) subsets from corresponding ground truth G c s, and (cid:98)
G pc
=
(cid:98) G ∗ c − (cid:98) G lc
and (cid:101) G pc = (cid:101) G ∗ c − (cid:101)
G lc be the complement, or equivalently, the p artial (cid:98) G ∗ c , (cid:101) G ∗ c that are estimated in (cid:98) G c , (cid:101) G c , respectively.
We can also define similar counterparts for G s : (cid:98) G ps , (cid:101) G ps are the partial (cid:98) G s , (cid:101) G s s contained in the estimated (cid:98) G c , (cid:101) G c while (cid:98) G ls , (cid:101)
G ls are the left subsets (cid:98) G s , (cid:101) G s , respectively.
Recall the constraint that | G c | =
s c , hence if (cid:98) G pc ⊆ (cid:98) G c , then a corresponding (cid:98) G lc = (cid:98) G ∗ c − (cid:98) G pc will be replaced by (cid:98) G ps in (cid:98) G c .
In this case, we have:
[Equation]
where the second equality is due to E = ˆ e is determined so that H ( E = ˆ e | (cid:98) G c , Y ) =
Compared
Eq.
19 to that when (cid:98) G c
=
(cid:98) G ∗ c , we have the entropy change as:
[Equation]
Let ϵ
= H ( (cid:98) G ps | (cid:98) G pc , E =
ˆ e, Y ) .
In a idealistic setting, when the noise of the generation process S := f spu ( Y, E ) in PIIF tends to be 0 , i.e., ϵ → 0 , S is determined conditioned on E, Y , hence G s and any subsets of G s are all determined.
Then, it suffices to know that in Eq.
20, H ( (cid:98) G ps | (cid:98) G pc , E =
ˆ e, Y )
=while H ( (cid:98) G l c | (cid:98) G p c , E = ˆ e, Y ) >
0 since (cid:98) G l c can not be determined when given (cid:98) G p c , E = ˆ e, Y .
Thus, when some subset from G s is included in (cid:98) G c , it will minimize H ( (cid:98) G c , E
= ˆ e | Y ) .
However in practice, it is usual that ϵ > 0 .
Therefore, in the next, we will show how ϵ =
H ( (cid:98) G ps | (cid:98) G pc , E = ˆ e, Y ) can be cancelled thus leading to a smaller H ( (cid:98) G c , E
= ˆ e | Y ) , by considering the second term H ( (cid:98) G c , E
= ˆ e | (cid:101) G c , E = ˜
e, Y ) .
As for H ( (cid:98) G c , E = ˆ e | (cid:101) G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possible cases into two:
(i) One of (cid:98) G c and (cid:101) G c contains some subset of G s ,
i
[Equation]
e., (cid:98) G c contains some (cid:98) G ps ⊆ (cid:98) G s ; (ii) Both (cid:98) G c and (cid:101) G c contain some (cid:98) G ps ⊆ (cid:98) G s and (cid:101) G ps ⊆ (cid:101) G s , respectively.
For (i), we have:
[Equation]
Thus, we can write the change of H ( (cid:98) G c , E
= ˆ e | (cid:101) G c , E = ˜
e, Y ) between (cid:98) G c
= (cid:98) G pc ∪ (cid:98) G ps and (cid:98) G c = (cid:98) G ∗ c as:
[Equation]
pc
[Equation]
Combing ∆ H ( (cid:98) G c , E
= ˆ e | Y ) , we have: ∆
[Equation]
G c , E = ˆ e ;
[Equation]
where the last equality is because of the independence of (cid:98) G ps between (cid:101) G c , E =
˜
e conditioned on
Y, E = ˆ e .
Since conditioning will lower the entropy for both discrete and continuous variables [22, 114], we have:
[Equation]
which implies the existence of (cid:98) G ps in (cid:98) G c will lower down the second term in Eq.
15 for the case (i).
For (ii), we have:
[Equation]
Similar to (i), H ( (cid:98) G ps | (cid:101) G pc , (cid:101) G ps , E = ˜
e, Y, (cid:98) G pc , E = ˆ e ) can be cancelled out with H ( (cid:98) G ps | (cid:98) G pc , E
= ˆ e, Y ) .
Then, we have:
[Equation]
Since additionally conditioning on (cid:98) G ps in H ( (cid:98) G lc , E = ˆ e | (cid:101) G pc , (cid:101) G ps , E = ˜
e, Y ) can not lead to new information about (cid:98) G lc , we have:
[Equation]
) which follows that ∆ I ( (cid:98) G c , E = ˆ e ; (cid:101) G c , E = ˜
e | Y ) < 0 .
To summarize, the ground truth G c is the only maximizer of the objective (Eq. 15), hence solving for the objective (Eq. 15) can elicit an invariant GNN.
E.3 Proof for theorem 3.1 (ii)
Theorem E.6 (CIGAv2 Induces Invariant GNNs) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c gen in 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] [Equation]
where (cid:98) G c
=
g ( G ) , (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G and G have the same label.
Similar to the proof for Theorem E.4, to prove Theorem E.6 is essentially to show the estimated (cid:98) G c through Eq.
28 is the underlying G c , hence the minimizer of Eq.
28 elicits an invariant GNN predictor (Definition.
E.1).
In the next, we also begin with a lemma:
Lemma E.7.
Given data generation process as Theorem E.6, for both FIIF and PIIF, we have:
I ( C ; Y ) ≥ I ( S ; Y ) , hence I ( G c ; Y ) ≥ I ( G s ; Y ) .
Proof for Lemma E.7.
For both FIIF and PIIF, Assumption 2.4 implies that H ( C | Y ) ≤ H ( S | Y ) .
It follows 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 c is invertible, we have I ( G c ; Y ) =
I ( C ; Y ) ≥ I ( S ; Y ) ≥ I ( G s ; Y ) .
Given Lemma E.7, we know (cid:98) G c at least contains some subset of the underlying G c , otherwise the constraint I ( (cid:98) G s ; Y ) ≤ I ( (cid:98) G c ; Y ) will be violated since G c ⊆ (cid:98) G s in this case.
Assuming there are some subset of G s contained in (cid:98) G c , without loss of generality, we can divide all of the possible cases about (cid:98) G c into two:
(i) (cid:98) G c only contains a subset of the underlying G c ; (ii) (cid:98) G c contains a subset of the underlying G c as well as part of the underlying G s ;
[Equation]
Before the discussion, let us inherit the notations of subsets of G c , G s from the proof for Theorem E.4:
Let (cid:98) G ∗ c and (cid:101) G ∗ c be the ground truth invariant subgraph G c s of (cid:98) G and (cid:101) G , (cid:98) G lc = (cid:98) G ∗ c − (cid:98) G c and (cid:101)
G lc = (cid:101) G ∗ c − (cid:101) G c be the l eft (un-estimated) subsets from corresponding ground truth G c s, and (cid:98)
G pc
=
(cid:98) G ∗ c − (cid:98) G lc
and (cid:101) G pc = (cid:101) G ∗ c − (cid:101)
G lc be the complement, or equivalently, the p artial (cid:98) G ∗ c , (cid:101) G ∗ c that are estimated in (cid:98) G c , (cid:101) G c , respectively.
Similarly, (cid:98) G ps , (cid:101) G ps are the partial (cid:98) G s , (cid:101) G s s contained in the estimated (cid:98) G c , (cid:101) G c while (cid:98) G ls , (cid:101) G ls are the left subsets (cid:98) G s , (cid:101) G s , respectively.
First of all, case (i) cannot hold because, when maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) , if ∃ (cid:98) G lc = (cid:98) G ∗ c − (cid:98) G c , as shown in the proof for Theorem E.4, including (cid:98) G lc into (cid:98) G c can always enlarge I ( (cid:98) G c ; (cid:101) G c | Y ) , while not affecting the optimality of I ( (cid:98) G s ; Y )
+
I ( (cid:98) G c ; Y ) by re-distributing (cid:98) G lc from (cid:98) G s to (cid:98) G c .
Consequently, (cid:98) G ∗ c must be included in (cid:98) G c , i.e., (cid:98) G ∗ c ⊆ (cid:98) G c .
As for case (ii), recall that, by the condition of equally distributed training samples from each training environment, maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) is essentially maximizing I ( (cid:98) G c , E
= ˆ e ; (cid:101) G c , E = ˜
e | Y ) , ∀ ˆ e, ˜ e ∈ E tr , hence, we have:
[Equation]
We claim Eq.
29 can eliminate any potential subsets in the estimated (cid:98) G c .
Similarly, we have:
[Equation]
where the second equality is due to E
= ˆ e is determined.
Compared to the case that (cid:98) G c
=
(cid:98) G ∗ c , we have:
[Equation]
Then, as for H ( (cid:98) G c , E = ˆ e | (cid:101) G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possible cases into two:
(a) (cid:98) G c contains some (cid:98) G ps ⊆ (cid:98) G s ; (b) Both (cid:98) G c and (cid:101) G c contain some (cid:98) G ps ⊆ (cid:98) G s and (cid:101) G ps ⊆ (cid:101) G s , respectively.
For (a), we have:
[Equation]
Similarly to the proof for Theorem E.4, when considering ∆ I ( (cid:98) G c ; (cid:101) G c | Y ) , the effects of
H ( (cid:98) G ps | (cid:101) G c , E = ˜
e, Y, (cid:98) G ∗ c , E = ˆ e ) is cancelled out by H ( (cid:98) G ps | (cid:98) G ∗ c , E
= ˆ e, Y ) .
Hence, we have:
[Equation]
For (b), we have:
[Equation]
Similarly, H ( (cid:98) G ps | (cid:101) G ∗ c , (cid:101) G ps , E = ˜
e, Y, (cid:98) G ∗ c , E = ˆ e ) = 0 can also be cancelled out by H ( (cid:98) G ps | (cid:98) G ∗ c , E
= ˆ e, Y ) .
Moreover, for H ( (cid:98) G ∗ c | (cid:101) G ∗ c , (cid:101) G ps , E = ˜
e, Y, E
= ˆ e ) , (cid:101) G ps can not bring no additional information about (cid:98) G ∗ c , when conditioning on (cid:101) G ∗ c , Y, E = ˜
e .
Hence, we also have:
[Equation]
To summarize, when maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) , including any (cid:98) G ps ⊆ (cid:98) G ∗ s can not bring additional benefit while affecting the optimality of I ( (cid:98) G s ; Y )
+ I ( (cid:98) G c ; Y ) .
More specifically, when considering the changes to I ( (cid:98) G s ; Y )
+ I ( (cid:98) G c ; Y ) ,
∀ G ps ⊆
G s , we have
[Equation]
while I ( Y ; (cid:98) G
[Equation]
c , G ps )
=
I ( Y ; (cid:98) G ∗ c )
+ I
( Y ; (cid:98) G ps | (cid:98) G ∗ c ) , ∀ e ∈ E tr .
Consequently,
[Equation]
[Equation]
Hence, only the underlying G c is the solution to Eq.
28, which implies that solving for the objective (Eq. 28) can elicit an invariant GNN.
F Details of Prototypical CIGA Implementation
In fact, the CIGA framework introduced in Sec.
3 can have multiple implementations.
We choose interpretable architectures in our experiments for the purpose of concept verification.
More sophisticated architectures can be incorporated.
Experimental results in Sec. 4 also demonstrates that, even equipped 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 CIGA decomposes a GNN model for graph classification into two modules, i.e., a featurizer:
g : G → G c and a classifier f c : G c → Y .
Specifically, for the implementation of Featurizer, we choose one of the common practices GAE [44] for calculating the sampled weights for each edge.
More formally, the soft mask is predicted through the following equation:
[Equation]
where a calculates the sampling weights for each edge using a MLP: M ij =
MLP ([ Z i , Z j ]) .
[Figure]
C Le A need to classify on the
The featurizer g will extract an G from each input for the classifier f c to predict the label.
The training objective of g is implemented in a contrastive strategy where the distribution of (cid:98) G c at the latent sphere will be optimized to maximize the intra-class mutual information.
With the identified invariant subgraph G c , the predictions made by classifier f c based on G c are invariant to distribution shifts;
If a sampling ratio s c is predetermined, we sample s c of total edges with the largest predicted weights as a soft estimation of (cid:98) G c .
Then, the estimated (cid:98) G c will be forwarded to the classifier f c for predicting the labels of the original graph.
Although Theorem E.4 assumes s c is known, in real applications we do not know the specific s c .
Hence, in experiments, we select s c according to the validation performance.
To thoroughly study the effects of I ( (cid:98) G s ; Y ) comparing to CIGAv1, we stick to using the same s c and sampling process for CIGAv2, while CIGAv2 essentially requires less specific knowledge about ground truth r c hence achieving better empirical performance.
Moreover, once the sampled edges are determined, the classifier GNN can take either the original feature of the input graph or the learned feature from the featurizer as the new node attributes for (cid:98) G c .
We select the architecture according to the validation performance from some random runs.
For the implementation of the information theoretic objectives, we will use CIGAv2 for elaboration while the implementation of CIGAv1 can be obtained via removing the third term from
CIGAv2.
Recall that CIGAv2 has the following formulation:
[Equation]
where (cid:98) G c
=
g ( G ) , (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G and G have the same label.
In Sec. 3.3, we introduce a contrastive approximation for I ( (cid:98) G c ; (cid:101) G c | Y ) :
[Equation]
c ;
[Equation]
where positive samples ( (cid:98) G c , (cid:101) G c ) are the extracted subgraphs of graphs that have the same label of
G , negative samples are those with different labels, P g ( G |Y =
Y ) is the pushforward distribution of P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label
Y , h (cid:98) G c , h (cid:101) G c , h G ic are the graph presentations of the estimated subgraphs, and ϕ is the similarity metric for the graph presentations.
As M → ∞ , Eq.
36 approximates I ( (cid:98) G c ; (cid:101) G c | Y ) which can be regarded as a non-parameteric resubstitution entropy estimator via the von Mises-Fisher kernel density [1, 41, 101].
While for the third term I ( (cid:98) G s ; Y ) given the constraint I ( (cid:98) G s ; Y ) ≤
I ( (cid:98) G c ; Y ) , a straightforward implementation is to imitate the hinge loss:
[Equation]
where N is the number of samples, I is a indicator function that outputs 1 when the interior condition is satisfied otherwise 0 , and R (cid:98) G s and R (cid:98) G c are the empirical risk vector of the predictions for each sample based on (cid:98) G s and (cid:98) G c respectively.
One can also formulate Eq.
35 from game-theoretic perspective [14].
Finally, we can derive the specific loss for the optimization of CIGAv2 combining Eq.
36 and Eq. 37:
[Equation]
where R (cid:98) G c , R (cid:98) G s are the empirical risk when using (cid:98) G c , (cid:98) G s to predict Y through the classifier.
Typically, we use a additional MLP downstream classifier
ρ s for (cid:98) G s in the classifier GNN.
h (cid:98) G c is the graph representation of (cid:98) G c which can be induced from the GNN encoder either in the featurizer or in the classifier.
α, β are the weights for I ( (cid:98) G c ; (cid:101) G c | Y ) and I ( (cid:98) G s ; Y ) , and ϕ is implemented as cosine similarity.
The optimization loss for CIGAv1 merely contains the first two terms in Eq.
The detailed algorithm for CIGA is given in the Algorithm 1, assuming the h (cid:98) G c is obtained via the graph encoder in f c .
Fig.
15 also shows a illustration of the working procedure of CIGA.
G Detailed Experimental Settings
In this section, we provide more details about our experimental settings in Sec. 4, including the dataset preparation, dataset statistics, implementations of baselines, selection of models and hyperparameters as 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 in our experiments.
Statistics of the datasets are presented in Table 4.
SPMotif datasets.
We construct 3-class synthetic datasets based on BAMotif [116, 58] following [104], where the model needs to tell which one of three motifs (House, Cycle, Crane) that the graph contains.
For each dataset, we generate 3000 graphs for each class at the training set,graphs for each class at the validation set and testing set, respectively.
During the construction, we merely inject the distribution shifts in the training data while keep the testing data and validation data without the biases.
For structure-level shifts ( SPMotif-Struc ), we introduce the bias based graph representations subgraphs contrastive positive samples and negative samples are constructed from the for the
[Table]
on FIIF, where the motif and one of the three base graphs (Tree, Ladder, Wheel) are artificially (spuriously) correlated with a probability of various biases, and equally correlated with the other two.
Specifically, given a predefined bias b , the probability of a specific motif (e.g., House) and a specific 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 influences of structure level shifts.
Moreover, to simulate more realistic scenarios where both structure level and topology level have distribution shifts, we also construct SPMotif-Mixed for mixed distribution shifts.
We additionally introduced FIIF attribute-level shifts based on SPMotif-Struc, where all of the node features are spuriously correlated with a probability of various biases by setting to the same number of corresponding labels.
Specifically, given a predefined bias b , the probability that all of the node features of a graph has label y
(e.g., y = 0 ) being set to y (e.g., X = 0 ) is b while for the others is (1 − b ) / 2 (e.g., P ( X = 1 )
=
P ( X = 2 ) =
(1 − b ) / 2 ).
More complex distribution shift mixes can 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.
[113], Bevilacqua et al.
[11] to study the OOD generalization abilities of various methods on four of TU datasets [67], i.e., PROTEINS, DD, NCI1, NCI109 .
Specifically, we use the data splits generated by Yehudai et al.
[113] and use the Matthews correlation coefficient as evaluation metric following [11] due to the class imbalance in the splits.
The splits are generated as follows:
Graphs with sizes smaller than the 50 -th percentile are assigned to training, while graphs with sizes larger than the 90 -th percentile are assigned to test.
A validation set for hyperparameters tuning consists of 10% held out examples from training.
We also provide a detailed statistics about these datasets in table 5.
[Table]
Graph-SST datasets.
Inspired by the data splits generation for studying distribution shifts on graph sizes, we split the data curated from sentiment graph data [122], that converts sentiment sentence classification datasets SST5 and SST-Twitter [90, 26] into graphs, where node features are generated using BERT [25] and the edges are parsed by a Biaffine parser [32].
Our splits are created according to the averaged degrees of each graph.
Specifically, we assign the graphs as follows:
Those that have smaller or equal than 50 -th percentile averaged degree are assigned into training, those that have averaged degree large than 50 -th percentile while smaller than 80 -th percentile are assigned to validation 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.
Specifically, the original Mnist dataset are assigned to binary labels where images with digits 0 − 4 are assigned to y = 0 and those with digits 5 − 9 are assigned to y = 1 .
Then, y will be flipped with a probability of 0 . 25 .
Thirdly, green and red colors will be respectively assigned to images with labels 0 and 1 an averaged probability of 0 .
15 (since we do not have environment splits) for the training data.
While for the validation and testing data the probability is flipped to 0 .
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 binding affinity prediction for both macromolecule (protein target) and small-molecule (drug compound).
The molecule data and the notations are curated from realistic ChEMBL database [63].
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 uses ic50 measurement type and contains core level annotation noises.
For more details, we refer interested readers to Ji et al.
G.2 Training and Optimization in Experiments
During the experiments, we do not tune the hyperparameters exhaustively while following the common recipes for optimizing GNNs.
Details are as follows.
GNN encoder.
For fair comparison, we use the same GNN architecture as graph encoders for all methods.
By default, we use 3 -layer GNN with Batch Normalization [39] between layers and JK residual connections at last layer [106].
For the architectures we use the GCN with mean readout [45] for all datasets except Proteins where we empirically observe better validation performance with a
GIN and max readout [107], and for DrugOOD datasets where we follow the backbone used in the paper [40], i.e., 4 -layer GIN with sum readout.
The hidden dimensions are fixed as 32 for SPMotif,
TU datasets, CMNIST-sp, and 128 for SST5, Twitter and DrugOOD datasets.
Optimization and model selection.
By default, we use Adam optimizer [43] with a learning rate of 1 e − 3 and a batch size of 32 for all models at all datasets.
Except for DrugOOD datasets, we use a batch size of 128 following the original paper [40].
To avoid underfitting, we pretrain models forepochs for all datasets, except for CMNIST and Twitter where we pretrain 5 epochs and for SST we 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 [91] 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 the validation 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 author released codes [120, 78], where we use the codes provided by the authors 6 for DIR c[104] which is the same as the author released codes.
During the implementation, we use the same s c for all interpretable GNN baselines, chosen from { 0 . 1 , 0 . 2 , 0 . 25 , 0 . 3 , 0 . 4 , 0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 . 9 } according to the validation performances, and set to 0 .
25 for SPMotif following Wu et al.
3 for Proteins and 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 can be unstable during its nested optimization for approximating the mutual information of the predicted subgraph and the input graph.
We use a larger batch size of 128 or reduce the nested optimization steps to be lower than 20 for stabilizing the performance.
If the optimization failed due to the instability during 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 by the author, because DIR adopts Last Epoch Model Selection which is different from the claim that they select models according to the validation performance , i.e., line 264 to linein train/spmotif dir.py from the commit 4b975f9b3962e7820d8449eb4abbb4cc30c1025d of https://github.com/Wuyxin/DIR-GNN .
We select the hyperparamter for the proposed DIR regularization from { 0 . 01 , 0 . 1 , 1 , 10 } according to the validation performances at the datasets, while we stick to the authors claimed hyperparameters for the datasets they also experimented with.
For invariant learning, we refer to the implementations in DomainBed [34] for IRM [4], V-Rex [49] and IB-IRM [2].
Since the environment information is not available, we perform random partitions on the training data to obtain two equally large environments for these objectives.
Moreover, we select the weights for the corresponding regularization from { 0 . 01 , 0 . 1 , 1 , 10 , 100 } for these objectives according to the validation performances of IRM and stick to it for others, since we empirically observe that they perform similarly with respect to the regularization weight choice.
For EIIL [23], we use the author released implementations about assigning different samples the weights for being put in each environment and calculating the IRM loss.
Besides, for CNC [124], we follow the algorithm description to modify the sampling strategy in supervised contrastive loss [42] based on a pretrained GNN optimized with ERM, and choose the weight for contrastive loss using the same grid search as for CIGA.
Implementations of CIGA.
For fair comparison, CIGA uses the same GNN architecture for GNN encoders as the baseline methods.
We did not do exhaustive hyperparameters tuning for the loss
Eq.
By default, we fix the temperature to be 1 in the contrastive loss, and merely search α from { 0 . 5 , 1 , 2 , 4 , 8 , 16 , 32 } and β from { 0 . 5 , 1 , 2 , 4 } according to the validation performances.
For
CMNIST-sp, we find larger β are required to get rid of intense spurious node features hence we expand 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 with few random runs.
Besides, we also have various implementation options for obtaining the features in (cid:98) G c , for obtaining h (cid:98) G c , as well as for obtaining predictions based on (cid:98) G s .
By default, we feed the graph representations of featurizer GNN to the classifier GNN, as well as to the contrastive loss.
For classifying G based on (cid:98) G s , we use a separate MLP downstream classifier in the classifier GNN f c .
The only exception is for the CMNIST-sp dataset where the spurious correlation is stronger than the 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 (cid:98) G s or (cid:98) G c , which we leave for future work.
Evaluation protocol.
We run each experiment 10 on TU datasets and 5 times for others where the random seeds start from 1 to the number of total repeated times.
During each run, we select the model according to the validation performance and report the mean and standard deviation of the corresponding metrics.
G.3 Software and Hardware
We implement our methods with PyTorch [73] and PyTorch Geometric [29].
We ran our experiments on 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 graphics cards with CUDA 10.2, 2 NVIDIA RTX 2080Ti and 2 NVIDIA RTX 3090Ti graphics cards with
CUDA 11.3, and NVIDIA TITAN series with CUDA 11.3.
G.4 Additional Analysis
Hyperparameter sensitivity analysis.
To examine how sensitive CIGA is to the hyperparamters α and β for contrastive loss and hinge loss, respectively, under different distribution shifts.
We conduct experiments based on the hardest datasets from each table (i.e., SPMotif-Mixed with the bias of 0 .
9 , DrugOOD-Scaffold and the NCI109 datasets from Table 1, Table 2, and Table 3, respectively.)
To increase the difficulty, we search for more fine-grained spaces for both parameters, i.e., { 0 . 1 , 0 . 5 , 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 } .
During changing the value of β , we will fix the α to a specific value under which the model has a relatively good performance (but not the best, to fully examine the robustness of CIGA in practice).
During the sensitivity tests, we follow the evaluation protocol as that used for the main experiments.
The results are shown in Fig.
16 and Fig.
[Figure] [Figure]
From the results above
[Figure]
we can see that both CIGAv1 and CIGAv2 are robust to different values of α and β , respectively, across different datasets and distribution shifts.
Notably, in Fig.
[Figure]
when the coefficient α for the contrastive loss become too small
[Figure]
the invariance of the identified invariant subgraphs (cid:98) G c may not be guaranteed, resulting worse performances.
Moreover, when α becomes too large, it may affect the optimization and yield worse performances.
In SPMotif datasets, the worse performances can be observed via the large variances as well.
Similarly for β , as shown in
Fig.
17, when β becomes too small, some part from the spurious subgraph may still be contained in the estimated invariant subgraphs.
While if β becomes too large, there might be part of (cid:98) G c being eliminated.
Although both CIGAv1 and CIGAv2 are robust to the changes of α and β , the intrinsic difficult optimization in OOD generalization algorithms including the proposed CIGA in our work, still require a more proper and smooth optimization process [18].
[Table]
Running time analysis.
To examine how much computational overhead is induced by the architecture and the additional objectives in CIGA, we analyze and compare the averaged training time of different methods on DrugOOD-Scaffold.
Factors that could affect the running time such as GNN backbone, batch size, and the running devices (NVIDIA RTX 2080Ti, Linux Servers with 40 cores Intel(R)
Xeon(R)
Silver 4114 CPU @
2.20GHz, 256 GB Memory, and Ubuntu 18.04 LTS), are fixed the same during the testing.
The results are shown as in Table.
It can be found that CIGA is the only OOD method that outperforms ERM by a non-trivial margin with a relatively low additional computational overhead.
[Table] [Table]
Single environment OOD generalization.
The theory of invariant learning fundamentally assume the presence of multiple environments [76, 4].
However in practice, it does not always hold, which would inevitably fail all of the invariant learning solutions [4, 49, 23, 2], including CIGA.
Nevertheless, to examine how CIGA performs under various realistic scenarios
[Table]
we conduct an additional experiment based on DrugOOD-Assay.
We select samples that are from the largest assay group (i
[Table]
e., the biochemical functionalities of these molecules are tested and reported under the same experimental setup in the lab)
The results are separated and shown in Table 7 and Table
[Table]
Besides the baselines, we also show the “Oracle” performances from the main table, to demonstrate the performance gaps.
From the Table 7 and Table
[Table]
we can see that, both CIGAv1 and CIGAv2 maintain their state-of-theart performances even in the single training environment setting.
We hypothesize that enforcing the mutual information between the estimated (cid:98) G c also helps to retain the invariance even under the single training environment setting.
That may partially explain why CNC can bring some improvements.
We believe it is an interesting and promising future direction to develop in-depth understanding and better solutions under this circumstance.
G.5 Interpretation Visualization
Since we use the interpretable GNN architecture to implement CIGA 7 , it brings an additional benefit that provides certain interpretation for the predictions automatically, which may facilitate human understanding in practice.
First, we provide some interpretation visualizations in SPMotif-Struc and SPMotif-Mixed datasets, under the biases of 0 .
6 and 0 . 9 .
Shown in Fig.
18 to Fig.
21, we use pink to color the ground truth nodes in G c , and denote the relative attention strength with edge color intensities.
Besides, we also provide some interpretation visualization examples in DrugOOD datasets.
Shown in
Fig.
22 to Fig.
27, we use the edge color intensities to denote the attentions of models that pay to the corresponding edge.
Some interesting patterns can be found in the molecules shared with the same label, which could provide insights to the domain experts when developing new drugs.
We believe that, because of its superior OOD generalization performance on graphs, CIGA can have high potential to push forward the developments of AI-Assisted Drug Discovery, and enrich the AI tools for facilitating the fundamental practice of science in the future.
[Figure] [Figure]
[Figure] [Figure]
[Figure]
C C F
N
C N
S O O
C C
C
C C
[Figure]
[Figure]
C C C
C N
[Figure] [Figure]
C
[Figure]
N
C C C C C C C
C O O
[Figure] [Figure]
[Figure] [Figure]
C
[Figure]
N
C C
C
[Figure]