Invariance Principle Meets 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 difﬁcult 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 to capture the invariance of graphs for guaranteed OOD generalization under various distribution shifts.
Speciﬁcally, 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 on both synthetic and real-world datasets, including a challenging setting in AI-aided drug discovery, validate the superior OOD generalization ability of our method.
1. Introduction
Graph neural networks (GNNs) have gained great success in tasks involving relational information (Kipf & Welling, 2017; Hamilton et al., 2017; Veliˇckovi´c et al., 2018; Xu et al., 2018; 2019).
However, it assumes that the training and test graphs are drawn from the same distribution, which is often violated in reality (Hu et al., 2020; Koh et al., 2021;
Huang et al., 2021; Ji et al., 2022).
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 (Beery et al., 2018;
DeGrave et al., 2021).
Such out-of-distribution (OOD) generalization failures become the major roadblock for practical applications of graph neural networks (Ji et al., 2022).
Meanwhile, enabling OOD generalization on Euclidean data has received surging attention and several solutions were proposed (Arjovsky et al., 2019;
Sagawa* et al., 2020; Bengio et al., 2020; Krueger et al., 2021; Creager et al., 2021;
Koyama & Yamaguchi, 2020; Ahuja et al., 2021).
In particular, the invariance principle from causality is at the heart of those works (Peters et al., 2016).
The principle leverages the
Independent Causal Mechanism (ICM) assumption (Pearl, 2009; Peters et al., 2017) and implies that, predictions that focus only on the causes of the label can stay invariant to a large class of distribution shifts (Peters et al., 2016).
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 can happen at both attribute-level and structure-level, and be observed in multiple forms (Wu et al., 2022a), where each shift can spuriously correlate with labels in different modes (Nagarajan et al., 2021; Ahuja et al., 2021).
The entangled distribution shifts make it more difﬁcult to identify the invariance.
Second, OOD methods developed on Euclidean data often require additional environment (or domain) labels for distinguishing the sources of distribution shifts (Arjovsky et al., 2019).
However, the environment labels could be highly expensive to obtain and often unavailable for graphs, as collecting the labels usually requires expert knowledge due to the abstraction of graphs (Hu et al., 2020).
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 a new framework,
G raph O ut-
Of -Distribution G eneralization (GOOD), to
[Figure]
enable guaranteed OOD generalization on graphs under different distribution shifts.
Speciﬁcally, we build three
Structural Causal Models (SCMs) (Pearl, 2009) to characterize the distribution shifts that could happen on graphs (Sec. 2.1).
Then, we generalize the invariance principle to graphs for OOD generalization: GNNs are invariant to distribution shifts if they focus only on a invariant and critical subgraph G c that contains the most of the information in
G about the underlying causes of the label.
Thus, we can achieve OOD generalization on graphs with two processes: invariant subgraph identiﬁcation and label prediction.
Accordingly, shown as Fig.
1, we implement a prototypical algorithm that decomposes a GNN into: a) a featurizer g for identifying the invariant subgraph G c from G ; b) a classiﬁer f c for making predictions based on G c .
To identify the desired G c , we derive an information-theoretic objective for g to extract subgraphs that maximally preserve the invariant intra-class information.
We show that this approach can provably identify the underlying G c (Sec. 3).
Experiments on 16 synthetic and realistic datasets with various distribution shifts, including a challenging setting from
AI-aided drug discovery (Ji et al., 2022), show that GOOD can signiﬁcantly outperform all of existing methods, demonstrating its promising OOD generalization ability (Sec. 4).
To our best knowledge, there is no existing work that could handle more comprehensive graph distribution shifts than
GOOD while also with OOD generalization guarantees.
Discussions on related works are deferred to Appendix C.2.
2. Graph OOD through the Lens of Causality
Problem Setup.
In this work, we focus on OOD generalization in graph classiﬁcation.
Speciﬁcally, we are given a set of graph datasets D =
{D e } e collected from multiple
[Figure]
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 classiﬁer ρ : 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 (Vapnik, 1991; Arjovsky et al., 2019).
We leave more background details in Appendix C.1.
2.1. Graph Generation Processes
It is known that OOD generalization is impossible without assumptions on the environments E all (Pearl, 2009; Ahuja et al., 2021).
In the next, we brief the assumptions on graph generation with SCMs and leave full details to Appendix A.
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.
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 (Kügelgen et al., 2021).
Similarly in images, C and S can represent content and style while E can refer to the locations where the images are taken (Zhang et al., 2021; Kügelgen et al., 2021).
C and
S further control the generation of the observed graphs and can have multiple types of interactions (Ahuja et al., 2021).
Graph generation model.
The SCM for graph generation is given as Fig. 2(a).
f gen is decomposed into three processes 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 (Pearl, 2009).
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 (Bohacek et al., 1996; Sterling & Irwin, 2015; Ji et al., 2022).
In contrast, the generation of G s and G will be affected by environment E G ⊆ E .
Thus, graphs collected from different environments (or domains) can have different structure-level (e.g., graph sizes (Bevilacqua et al., 2021)) as Training Data Testing Data “House” “House” “Cycle” ( [Figure]
0.9 Data Biases
A cc u r a c y
Struc-ERMMixed-ERMStruc-IRM Mixed-IRM (b) Structure and attribute shifts. 0.9 Data Biases A cc u r a c y
Struc-ERMMixed-ERMStruc-IRM
Mixed-IRM
[Figure]
{+well as feature-level statistics (e.g., homophily (McPherson et al., 
+}
; Chen et al., 2022; Wu et al., 2022a)).
Interactions at latent space.
Following previous works (Arjovsky et al., 2019; Ahuja et al., 2021), we categorize the interactions between C and S into Fully Informative Invariant Features (FIIF, Fig. 2(b)) and Partially Informative
Invariant Features (PIIF, Fig. 2(c)), depending on whether C is fully informative about Y , i.e., ( S, E ) ⊥⊥ Y | C .
In the two
SCMs, 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 (Aubin et al., 2021; Nagarajan et al., 2021).
This issue can be more serious in graphs, since different distribution shifts can have different interaction modes at the latent space.
Moreover, 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 (Burshtein et al., 1992;
Chapelle et al., 2006; Schölkopf, 2019; Schölkopf et al., 2021), which also serves as the necessary separation assumption for a classiﬁcation task (Muller et al., 2001; Chen et al., 2005; Mika et al., 1999), i.e., H ( C | Y ) ≤ H ( S | Y ) .
2.2. Challenges of OOD Generalization on Graphs
Built upon the graph generation process, can existing methods produce a desired invariant GNN model?
Using the
BAMotif task (Luo et al., 2020) as Fig.
3, we ﬁnd that, neither ERM nor IRM, or more expressive GNN architectures can help improve the OOD generalization ability of GNNs.
The main reasons are: 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.
More results are given in Appendix D.
Invariance Principle for Graph OOD
Aiming to bridge the gap, we propose GOOD:
G raph O ut-
Of -Distribution G eneralization, to generalize and instantiate the invariance principle on graphs.
Full details and theoretical analysis are deferred to Appendix B.
Invariance for OOD generalization on graphs.
According to the ICM assumption (Peters et al., 2017), the labeling process C → Y in Fig.
2 is not informed nor inﬂuenced by other processes, implying that the conditional distribution
P ( Y | C ) remains invariant to the interventions on the environment latent variable E (Pearl, 2009).
Consequently, for a
GNN with a permutation invariant encoder h : G → R h and a downstream classiﬁer ρ : 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 (Vapnik, 1991) and can stay invariant to distribution shifts.
Causal algorithmic alignment.
To enable a GNN to learn to extract the information about C from G , we propose to explicitly aligns with the two causal mechanisms during the graph generation, i.e., C → G and ( G s , E G , G c ) → G , motivated by Xu et al.
Speciﬁcally, we realize the alignment by decomposing a GNN into two sub-components: a) a featurizer GNN g : G → G c aiming to identify the desired
G c ; b) a classiﬁer 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 R ( f c ( ˆ G c )) is the empirical risk of f c that takes ˆ G c as inputs to predict label Y for G , and ˆ G c is the intermediate subgraph containing information about C and hence needs to be independent of E .
Moreover, the extracted G c can either share the same graph space with input G or has its own
[Table]
space with latent node and edge features, depending on the speciﬁc graph generation process.
In practice, interpretable
GNN architectures (Yuan et al., 2020) are compatible with
GOOD, hence can serve as practical choices for GOOD.
Details are given in Appendix F.
Optimization objective.
To ensure the independence constraint ˆ G c ⊥⊥ E under the absence of E , we translate other properties of G c from SCMs in Sec.
2.1 into differentiable and equivalent conditions.
In a simplistic setting where all the invariant subgraphs G c have the same size s c , i.e., | G c | =
s c .
We derive the ﬁrst objective (GOODv1):
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
In Theorem B.2, we discuss why Eq. 2 is equivalent to Eq.
Although being effective,
Eq.
2 requires a strong assumption about the size of G c .
However, the size of G c is usually unknown or changes for different C s. In this circumstance, maximizing Eq.without additional constraints would lead to the presence of subgraphs of G s in ˆ G c .
For instance, ˆ G c =
G is a trivial solution to Eq.
2 when s c = ∞ .
To circumvent this limitation, we further resort to the properties of G s and obtain a new objective GOODv2 as follows
[Equation]
where ˆ G c , ˜ G c and ˜ G are the same as Eq.
We deffer the theoretical analysis and implementation details to Appendix B.
4. Empirical Studies
We conduct extensive experiments with 16 datasets to verify the effectiveness of GOOD.
We give analysis on the synthetic datasets and more details to Appendix G.
Datasets.
We use the SPMotif datasets from DIR (Wu et al., 2022c) where artiﬁcial structural shifts and graph size shifts are nested (SPMotif-Struc).
Besides, we construct a harder version mixed with attribute shifts (SPMotif-Mixed).
Baselines and our methods.
Besides ERM, we also compare with SOTA interpretable GNNs, GIB (Yu et al., 2021),
ASAP Pooling (Ranjan et al., 2020), and DIR (Wu et al., 2022c), to validate the effectiveness of the optimization objective in GOOD.
To validate the effectiveness of the decomposition in GOOD, we compare GOOD with SOTA
OOD objectives including IRM (Arjovsky et al., 2019), vRex (Krueger et al., 2021) and IB-IRM (Ahuja et al., 2021),
EIIL (Creager et al., 2021) and CNC (Zhang et al., 2022).
More comparison details are deferred to Appendix G.3.
Evaluation.
We report the mean and standard deviation of classiﬁcation accuracy for all datasets from multiple times.
OOD 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 converge 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.2 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,
GOODv1 and GOODv2 outperform all of the baselines by a signiﬁcant margin up to 10% with lower variance , which demonstrates the effectiveness and excellent OOD generalization ability of GOOD.
More analysis and results on real-world datasets are given in Appendix G.
Conclusion.
We studied the OOD generalization on graphs via graph classiﬁcation, and propose a new solution
GOOD 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. Acknowledgements
We thank the reviewers for their valuable comments.
This work was supported by GRF 14208318 from the RGC of
HKSAR and CUHK direct grant
E Theory and Discussions
E.2 Proof for theorem B.2 (i) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.3 Proof for theorem B.2 (ii) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Details of Prototypical GOOD Implementation
G Full Experimental Results and Settings
G. 1 Full experimental results and analysis
G.2 Details about the datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . G.
3 Training and Optimization in Experiments
G.
4 Software and hardware
A. Full Structural Causal Models on Graph Generation
Supplementary to the graph generation process in Sec.
2.1, we provide full SCMs on the graph generation process in this section as shown in Fig. 4.
Formal descriptions are given as Assumptions A.1, A.2, A.3, A.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 (Kügelgen et al., 2021; Ahuja et al., 2021), 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 (Beery et al., 2018;
Zhang et al., 2021; Kügelgen et al., 2021).
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 (Scaffold), experimental environment for examining the chemical property (Assay) (Ji et al., 2022).
On the contrary, the other latent variable S inherits environment-speciﬁc information thus can further affect the ﬁnally 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 (Ahuja et al., 2021).
Assumption A.1 (Graph Generation SCM) .
[Equation]
Speciﬁcally, the graph generation process is shown as Fig. 4(a).
The generative 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 c X ) , G c , ( Z sA , Z s X ) , 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 cX and Z sX 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 (Snijders & Nowicki, 1997; Lovász & Szegedy, 2006; You et al., 2018;
Luo et al., 2021; Bevilacqua et al., 2021).
Note that since our focus is to describe the potential distribution shifts with SCMs, in Assumption A.1, we aim to build a SCM that is compatible to many graph generation processes (Snijders &
Nowicki, 1997; Lovász & Szegedy, 2006; You et al., 2018; Luo et al., 2021), and leave speciﬁc graph generation processes and their implications to OOD generalization to future work.
Moreover, a subset of environment latent variable E G ⊆
E will affect the generation of G and G s .
Thus, graphs collected from different environments can have different structure-level properties such as degrees, graph sizes, and subgraph densities, as well as feature-level properties such as homophily (Knyazev et al., 2019; Yehudai et al., 2021; Bevilacqua et al., 2021;
Chen et al., 2022).
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 (Yu et al., 2021) and DIR (Wu et al., 2022c) mainly focus on FIIF case, while others such as IRM (Arjovsky et al., 2019) mainly focuses on the PIIF case.
Evidences show that failing to model either of them when developing the OOD objectives can have serious performance degenerations in practice (Aubin et al., 2021; Nagarajan et al., 2021).
That is why we aim to model both of them in our solution.
Assumption A.2 (FIIF SCM) .
[Equation]
Assumption A.3 (PIIF SCM) .
[Equation]
Assumption A.4 (MIIF SCM) .
[Equation]
As for the interactions between C and S at the latent space, we categorize the interaction modes into Fully Informative
Invariant Features (FIIF, Fig. 4(b)), and Partially Informative Invariant Features (PIIF, Fig. 4(c)), depending on whether the latent invariant part C is fully informative about label Y , i.e., ( S, E ) ⊥⊥ Y | C .
It is also possible that FIIF and PIIF are entangled into a Mixed Informative Invariant Features (MIIF,Fig. 4(d)).
We follow Arjovsky et al.
; Ahuja et al.
(2021) to formulate the SCMs for FIIF and PIIF, where we omit noises for simplicity (Pearl, 2009; Peters et al., 2017).
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
[Figure]
Among all of the interaction modes, f gen corresponds to the graph generation process in Assumption A.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 (Ahuja et al., 2021; Nagarajan et al., 2021).
Additionally, in MIIF, S is further partitioned into S 1 and 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 (Burshtein et al., 1992;
Chapelle et al., 2006; Schölkopf, 2019; Schölkopf et al., 2021), which also serves as the necessary separation assumption for a classiﬁcation task (Muller et al., 2001; Chen et al., 2005; Mika et al., 1999).
Assumption A.5 (Better Clustered Invariant Features) .
H ( C | Y ) ≤ H ( S | Y ) .
B. Full GOOD Framework B.1. Invariant Graph Neural Networks
To start, we formulate the desired GNN that is able to generalize to OOD graphs under different distribution shifts as below.
Deﬁnition B.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.1, considering a GNN ρ ◦ h that has a permutation invariant graph encoder h : G → R h and a downstream classiﬁer ρ : 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 .
B.2. Invariance for OOD Generalization on Graphs
Towards extending the invariance principle to graphs under SCMs in Sec. 2.1, we need to identify a set of variables that have stable causal relationship with Y under both FIIF and PIIF (Assumption A.2, A.3).
According to the ICM assumption (Peters et al., 2017), the labeling process C → Y is not informed nor inﬂuenced by other processes, implying that the conditional distribution P ( Y | C ) remains invariant to the interventions on the environment latent variable E (Pearl, 2009).
Consequently, for a GNN with a permutation invariant encoder h : G → R h and a downstream classiﬁer ρ : 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 (Vapnik, 1991) and can achieve the desired min-max optimality required by an invariant GNN (Def. B.1).
However, recovering
C from G is particularly difﬁcult, since the generation of G from C involves two causal mechanisms f G c gen and f G gen in
Assumption A.1.
The unavailability of E further adds up the difﬁculty of enforcing the independence between the learned representations and E .
B.3. Invariant Graph Learning Framework
Causal algorithmic alignment.
To enable a GNN to learn to extract the information about C from G , we propose the
GOOD framework that explicitly aligns with the two causal mechanisms f G c gen and f G gen in Assumption A.1.
The idea of alignment in GOOD 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 (Xu et al., 2020; 2021b).
Speciﬁcally, we realize the alignment by decomposing a GNN into two sub-components 1 : a) a featurizer GNN g : G → G c aiming to identify the desired G c ; b) a classiﬁer 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 R ( f c ( ˆ G c )) is the empirical risk of f c that takes ˆ G c as inputs to predict label Y for G , and ˆ G c is the intermediate subgraph containing information about C and hence needs to be independent of E .
Moreover, the extracted G c can either share the same graph space with input G or has its own space with latent node and edge features, depending on the speciﬁc graph generation process.
In practice, architectures from the literature of interpretable GNNs are compatible with
GOOD (Yuan et al., 2020), hence can serve as practical choices for the implementation of GOOD.
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 can not satisfy the condition ˆ G c ⊥⊥ E .
Formally, merely minimizing R ( f c ( ˆ G c )) is equivalent to maximizing a variational lower bound of I ( ˆ G c ; Y ) (Alemi et al., 2017; Yu et al., 2021), which may include a subgraph from G s in ˆ 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 objectives (Arjovsky et al., 2019; Krueger et al., 2021;
Ahuja et al., 2021; Sagawa* et al., 2020; Ganin et al., 2016; Sun & Saenko, 2016; Dou et al., 2019; Mahajan et al., 2021), making the identiﬁcation of G c more challenging.
Optimization objective.
To mitigate this issue, we need to translate other properties of G c from SCMs in Sec.
2.1 into differentiable and equivalent conditions to the independence constraint ˆ G c ⊥⊥ E .
To start, consider a simplistic setting where all the invariant subgraphs G c have the same size s c , i.e., | G c | =
s c 2 .
When maximizing I ( ˆ G c ; Y ) by minimizing R ( f c ( ˆ G c )) in Eq.
4, we are trying to extract all of the informative parts in G about Y into ˆ G c .
For 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 ( ˆ G c ; Y ) .
However, some subgraph of G c can be replaced by some subgraph of G s that is equally informative about
Y .
For PIIF (Fig. 2(c)), there also exists some subgraph of G s that contains additional information about Y than G c , hence ˆ
G c is more likely to involve some subgraph of G s .
Thus, the new condition needs to eliminate the auxiliary subgraphs of ˆ G c from G s such that the estimated ˆ G c can only contain G c .
Recall that for both FIIF and PIIF SCMs (Fig. 2), given two environments e 1 and e 2 : if G c appears in both e 1 and e 2 , the correctly identiﬁed ˆ G e 1 c , ˆ G e 2 c in e 1 and e 2 tend to have higher mutual information about the other, i.e., ( G c , G c ) ∈ arg max I ( ˆ G e 1 c ; ˆ G e 2 c ) .
While for G c and another G c (cid:48) corresponding to a different C (cid:48) (cid:54) =
C , if they appear in the same environment, then including any subgraph from G s in the estimated ˆ G c , ˆ G c (cid:48) for G c , G c (cid:48) will enlarge their mutual information, i.e., ( G c , G c (cid:48) )
∈ arg min I ( ˆ G c ; ˆ G c (cid:48) ) .
Thus, we can derive another important property about G c :
[Equation]
where ˆ G c and ˜ G c share the same C while ˆ G c (cid:48) corresponds to a different c (cid:48) .
In practice, we are also not given C .
However, since C and Y shares a stable causal relationship in both FIIF and PIIF SCMs, Y can serve as an alternative of C in
Eq. 5.
Moreover, when both I ( ˆ G c ; ˜ G c | C = c ) and I ( ˆ G c ; Y ) are maximized, I ( ˆ G c ; ˆ G c (cid:48)
| C = c (cid:48) , c (cid:48) (cid:54) =
c ) is automatically minimized, otherwise all classes will collapse to trivial solutions which is not possible given I ( ˆ G c ; Y ) being maximized.
Thus, we can replace the independence condition in Eq.
4 and obtain the following objective (GOODv1):
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
In Theorem B.2, we discuss why Eq.is equivalent to Eq. 4.
Although being effective, Eq. 6 requires a strong assumption about the size of G c .
However, the size of G c is usually unknown or changes for different C s. In this circumstance, maximizing Eq.
6 without additional constraints would lead to the presence of subgraphs of G s in ˆ G c .
For instance, ˆ G c =
G is a trivial solution to Eq.
6 when s c = ∞ .
To circumvent this limitation, we further resort to the properties of G s .
In both FIIF and PIIF SCMs (Fig. 2), G s and G c share certain overlapped information about Y via some subgraphs.
When maximizing I ( ˆ G c ; ˜ G c | Y ) and I ( ˆ G c ; Y ) , the appearances of these subgraphs in ˆ G c will not affect the optimality.
On the other hand, it can reduce the mutual information between the left part ˆ G s =
G − ˆ G c and Y , i.e., I ( ˆ G s ; Y ) .
In other words, by maximizing I ( ˆ G s ; Y ) , we can avoid involving additional subgraphs from G s into ˆ G c .
Meanwhile, to avoid trivial solution that G c ⊆
ˆ G s
during maximizing I ( ˆ G s ; Y ) , we can leverage the better clustering property of G c implied by Assumption A.5 to derive the constraint I
( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) .
Thus, we can obtain a new objective GOODv2 as follows:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
We also prove the equivalence between
Eq.
7 and Eq. 4 in Theorem B.2.
B.4. Theoretical Analysis and Practical Discussions
Theorem B.2 (GOOD 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.1, assuming that (a) f G gen and f G c gen in Assumption A.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 6, elicits an invariant GNN (Def. B.
(ii).
Each solution to Eq.
7, elicits an invariant GNN (Def. B
We prove Theorem B.2 (i) and (ii) in Appendix E.2, E.3, respectively.
Practical implementations of GOOD objectives.
After showing the power of GOOD, we introduce the practical implementations of GOODv1 and GOODv2 objectives.
Speciﬁcally, an exact estimate of the second term
I
( ˆ G c ; ˜ G c | Y ) could be highly expensive (van den Oord et al., 2018; Belghazi et al., 2018).
However, contrastive learning with supervised sampling provides a practical solution for the approximation (Khosla et al., 2020; Chopra et al., 2005; Salakhutdinov &
Hinton, 2007; van den Oord et al., 2018; Belghazi et al., 2018):
[Equation]
where positive samples ( ˆ G c , ˜ G c ) are the extracted subgraphs of graphs that have the same label of G , negative samples are those with different labels, P g ( G |Y
=
Y ) is the push-forward 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 (cid:54) =
Y ) refers to the distribution of G given the label that is different from Y , h ˆ G c , h ˜ G c , h G ic are the graph presentations of the estimated subgraphs, and φ is the similarity metric for graph presentations.
As M → ∞ , Eq. 8 approximates I ( ˆ G c ; ˜ G c | Y ) , which can be regarded as a non-parameteric resubstitution entropy estimator via the von Mises-Fisher kernel density (Ahmad & Lin, 1976; Kandasamy et al., 2015; Wang & Isola, 2020).
Thus, plugging it into Eq.
6 and Eq.
7 can relieve the issue of approximating I ( ˆ G c ; ˜ G c | Y ) in practice.
For the implementation of I ( ˆ G s ; Y ) and the constraint I ( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) in GOODv2, a practical choice is to follow the idea of hinge loss, I ( ˆ G s ; Y )
≈ 1 N R
ˆ G s · I ( R ˆ G s ≤ R ˆ G c ) , where N is the number of samples, I is an indicator function that outputs 1 when the inner condition is satisﬁed otherwise 0 , and R ˆ G s and R ˆ G c are the empirical risk vector of the predictions for each sample based on the corresponding ˆ G s and ˆ G c .
More implementation details can be found in Appendix F.
Discussions and implications of GOOD.
Although using contrastive learning to improve OOD generalization is not new in the literature (Dou et al., 2019; Mahajan et al., 2021; Zhang et al., 2022), previous methods cannot yield OOD guarantees in graph circumstances due to the highly non-linearity and the unavailability of domain labels E .
In particular, GOOD 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 (Zhang et al., 2022), yields unsatisfactory OOD generalization performance, which further implies the necessity of the decomposition in GOOD.
Moreover, the architecture of GOOD can have multiple other implementations for both the featurizer and classiﬁer, such as identifying G c at the latent space (Schölkopf, 2019; Schölkopf et al., 2021).
Since we cannot enumerate every possible implementation, in this work we choose interpretable GNN architectures as a prototype validation for GOOD and leave more sophisticated architectures as future works.
In particular, when optimized with ERM objective, GOOD can be reduced to interpretable GNNs .
However, merely using interpretable GNNs such as ASAP (Ranjan et al., 2020), GIB (Yu et al., 2021) or DIR (Wu et al., 2022c) cannot yield satisfactory OOD performance.
As shown in Table 2 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 C.
C. More Discussions on Related Works and Future Directions C.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 classiﬁcation, 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
= ρ ( r G ) with a downstream classiﬁer ρ : 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) (Xu et al., 2019;
Ying et al., 2018; Murphy et al., 2019; Xu et al., 2019; Chen et al., 2020; Morris et al., 2021), 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]
11) where R e is the empirical risk under environment
e (Vapnik, 1991; Peters et al., 2016; Arjovsky et al., 2019).
More details can be referred in (Ahuja et al., 2021).
C.2. Related work
We provide detailed related work discussion in this section in complementary to that in Introduction (Sec. 1).
To begin with, we summarize the main differences between our solution and them in Table 2.
An overview of related works.
On Euclidean data, Invariant Learning (Arjovsky et al., 2019; Creager et al., 2021; Ahuja et al., 2021), Group Distributionally Robust Optimization (GroupDro) (Krueger et al., 2021; Sagawa* et al., 2020;
Zhang et al., 2022), Domain Adaption (DA) and Domain Generalization (DG) (Ganin et al., 2016; Sun & Saenko, 2016;
Li et al., 2018b; Dou et al., 2019; Mahajan et al., 2021; Wang et al., 2021) are three widely adopted approaches to enable OOD generalization.
However, they all have limitations when being applied to graphs.
First, previous invariant learning methods are mostly developed and analyzed for Euclidean data (Arjovsky et al., 2019; Ahuja et al., 2021; Creager et al., 2021), or under speciﬁc SCM assumptions (Arjovsky et al., 2019), making the theoretical results hardly applicable to the complicated graph data (Rosenfeld et al., 2021) that can have multiple types of distribution shifts (Nagarajan et al., 2021).
GroupDro that minimizes the gap between worst group risk and average risk (Krueger et al., 2021; Sagawa* et al., 2020; Zhang et al., 2022), and DA/DG methods that aim to learn class-conditional domain invariant representations (Ganin et al., 2016; Sun & Saenko, 2016; Li et al., 2018b; Dou et al., 2019; Wang et al., 2021), cannot guarantee a min-max optimal predictor without additional assumptions (Arjovsky et al., 2019; Gulrajani & Lopez-Paz, 2021; Ahuja et al., 2021).
Moreover, most existing methods require environment labels that are however expensive to obtain in graphs, which limits their applications to graphs (Arjovsky et al., 2019; Krueger et al., 2021;
Ahuja et al., 2021; Sagawa* et al., 2020; Ganin et al., 2016; Sun & Saenko, 2016; Dou et al., 2019; Mahajan et al., 2021).
In contrast, we aim at a uniﬁed framework that are provably generalizable under different types of distribution shifts on graphs.
Another line of relevant works is about GNN explainability that aims to ﬁnd a subgraph of the input as the explanation for a GNN prediction (Ying et al., 2019; Yuan et al., 2020).
Although some may leverage causality to justify the generated explanation (Lin et al., 2021), they mostly focus on understanding the predictions of GNNs instead of for OOD generalization.
The most close works to ours are two interpretable GNNs that aim to explicitly extract a subgraph for both predictions and explanations.
However, they focus on graphs and shifts generated under a speciﬁc SCM.
Although one of them can provide theoretical guarantee for OOD generalization (Yu et al., 2021) by using the information bottleneck criteria (Ahuja et al., 2021), they would inevitably fail to generalize to graphs generated under different SCMs.
More discussions about the failure are given in Appendix D.4.
Besides, Bevilacqua et al.
(2021) also discuss OOD generalization on graphs but limited a speciﬁc graph family and graph size shifts.
Wu et al.
(2022b)
propose OOD methods on graphs for the task of node classiﬁcation but limited to shifts under a speciﬁc SCM.
Causality and OOD Generalization.
Causality comes to the stage for demystifying and improving the huge success of machine learning algorithms to further advances (Pearl, 2019; Schölkopf, 2019; Schölkopf et al., 2021).
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 inﬂuence the other conditional distributions (Pearl, 2009; Peters et al., 2017).
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 ﬁnding an invariant representation about the causal relationship between the underlying causes and the label (Peters et al., 2016; Arjovsky et al., 2019).
Consequently, models built upon the invariant representation can generalize to unseen environments or domains with guaranteed performance (Peters et al., 2016;
Rojas-Carulla et al., 2018;
Arjovsky et al., 2019; Sagawa* et al., 2020; Bengio et al., 2020; Koyama & Yamaguchi, 2020;
Gulrajani & Lopez-Paz, 2021; Krueger et al., 2021; Creager et al., 2021; Ahuja et al., 2021).
The arguably ﬁrst formulation of invariance principle was introduced by Peters et al.
Arjovsky et al.
(2019) 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 (Namkoong & Duchi, 2016) and generalization (Zhang et al., 2017), and prove its effectiveness in addressing
PIIF spurious correlations (Assumption A.3).
However, in practice, both PIIF and FIIF (Assumption A.2) can appear in data, while IRM can fail in these cases (Aubin et al., 2021; Nagarajan et al., 2021).
Ahuja et al.
(2021) 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.
(2021) and Lin et al.
(2022) propose new OOD objectives to relieve the needs for environment partitions, but limited to PIIF spurious types and linear regime.
In parallel invariant learning approaches, Sagawa* et al.
(2020) propose to regularize the worst group in group distributionally robust optimization (GroupDro).
Zhang et al.
(2022) 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(Rojas-Carulla et al., 2018; Chuang et al., 2020; Ganin et al., 2016; Sun & Saenko, 2016; Li et al., 2018b; Dou et al., 2019; Mahajan et al., 2021; Wang et al., 2021), 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 (Gulrajani & Lopez-Paz, 2021;
Ahuja et al., 2021), hence are not able to yield predictors with OOD generalization guarantees.
We refer interested readers to Pearl (2019); Schölkopf (2019); Schölkopf et al.
(2021) for an in-depth understanding, and Gulrajani & Lopez-Paz (2021);
Ahuja et al.
(2021) for a thorough overview.
GNN Explainability.
Works in GNN explainability aim to ﬁnd a subgraph of the input graph as the explanation for the prediction of a GNN model (Ying et al., 2019; Yuan et al., 2020).
Although some may leverage causality in explanation generation (Lin et al., 2021), 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 (Yu et al., 2021) and DIR (Wu et al., 2022c), and both of them focus on tackling FIIF spurious correlations (Assumption A.2).
The theoretical guarantees of GIB follows the theory of information bottleneck (Tishby et al., 1999), while GIB can not solve PIIF spurious correlations (Assumption A.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 (Arjovsky et al., 2019;
[Table]
Ahuja et al., 2021; Aubin et al., 2021; Nagarajan et al., 2021).
While for DIR, though as a generalization of Chang et al.
(2020) 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 (Santoro et al., 2018; Saxton et al., 2019), physics (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018) [Table] and graph algorithms (Tang et al., 2020; Velickovic et al., 2020; Xu et al., 2020;
Xhonneux et al., 2021).
Xu et al.
(2021b) study the neural network extrapolation ability from a geometrical perspective.
Han et al.
(2021) improve OOD drug discovery by mitigating the overconﬁdent misprediction issue.
Knyazev et al.
Yehudai et al.
(2021) 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.
(2021) study the graph size extrapolation problem of GNNs through a causal lens, while the induced invariance principle is built upon assumptions on the speciﬁc 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.
(2022b)
propose causal models as well as specialized objectives to extrapolate nodes with different neighbors.
However, their formulation is limited to node classiﬁcation task and speciﬁc spurious correlation type.
In contrast, the induced invariance principle in Wu et al.
(2022b), can be seen as a extension of GOOD for node classiﬁcation, 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 speciﬁc formulation and implementation to future works.
C.3. More discussions on connections of GOOD with existing work
Although primarily serving for graph OOD generalization problem, our theory complements the identiﬁability 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 ) (Zimmermann et al., 2021; Kügelgen et al., 2021).
Moreover, our results also partially explain the success of graph contrastive learning (You et al., 2020;
Ma et al., 2021;
You et al., 2021), where GNNs may implicitly learn to identify the underlying invariant subgraphs for prediction.
On expressivity of graph encoder in GOOD.
The expressivity of GOOD 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 classiﬁer 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 GOOD can improve the OOD generalization.
On GOOD 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 (Wu et al., 2022c) as they also focus on FIIF), which further implies Y ⊥⊥ S | ˆ G c .
Hence it is natural to use Information Bottleneck (IB) objective (Tishby et al., 1999) to solve for G c :
[Equation] [Equation]
which explains the success of many existing works in ﬁnding predictive subgraph through IB (Yu et al., 2021).
However, the estimation of I ( ˆ G c , G ) is notoriously difﬁcult due to the complexity of graph, which can lead to unstable convergence as observed in our experiments.
In contrast, optimization with contrastive objective in GOOD as Eq.
8 induces more stable convergence.
On GOOD for node classiﬁcations.
As the task of node classiﬁcation can be viewed as graph classiﬁcation based on the ego-graphs of a node, our analysis and discoveries can generalize to node classiﬁcation.
More speciﬁcally, the invariance principle for node classiﬁcation 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 (Wu et al., 2022b).
D. More Details about Failure Case Studies in Sec. 2.
In this section, we provide details on failure case studies in Sec. 2.2.
We ﬁrst 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.
Training Data
Testing Data
[Figure]
To begin with, we construct 3-class synthetic datasets based on BAMotif (Luo et al., 2020) and follow Wu et al. (2022c) to inject spurious correlations between motif graph and base graph during the generation.
In this graph classiﬁcation 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 artiﬁcial 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 (Kipf & Welling, 2017) with mean readout, a hidden dimension of 64 , and JK jump connections (Xu et al., 2018) at the last layer.
During training, we use a batch size of 32 , learning rate of 1 e − 3 with Adam optimizer (Kingma & Ba, 2015), and batch normalization between hidden layers (Ioffe & Szegedy, 2015).
Meanwhile, to stabilize the training, we also use dropout (Srivastava et al., 2014) of 0 .
1 and early stop the training when the validation accuracy does not increase till 5 epoch after ﬁrst 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 (Arjovsky et al., 2019), 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.
In Fig. 6, 7, 8, 9, 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.2.
More concretely, we focus on answering the following question from multiple aspects.
Can better architectures improve OOD generalization of GNNs?
Training with ERM objective.
As shown in Fig.
6, 7, 8, 9, we ﬁnd that GNNs trained with the standard empirical risk minimization (ERM) algorithm (Vapnik, 1991) 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 overﬁt to the shortcuts or spurious correlations presented in speciﬁc substructures or attributes in the graphs (Geirhos et al., 2020).
This phenomenon has also been shown to exist in GNNs equipped with more sophisticated architectures such as attention mechanisms (Veliˇckovi´c et al., 2018), under graph size shifts (Knyazev et al., 2019).
Training with OOD objective.
Meanwhile, as shown in Fig. 6, 7, 8, 9, OOD objectives primarily developed on Euclidean data such as invariant risk minimization (IRM) (Arjovsky et al., 2019) also cannot alleviate the problem.
On the contrary, IRM can fail catastrophically at non-linear regime if without sufﬁcient support overlap for the test environments, i.e., ∪ e ∈E te supp ( P e ) (cid:54)⊆ ∪ e ∈E tr supp ( P e ) (Rosenfeld et al., 2021).
In addition to IRM, the failure would also happen for alternative objectives (Krueger et al., 2021; Bellot & van der Schaar, 2020; Ahuja et al., 2021) 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 (Aubin et al., 2021; Nagarajan et al., 2021).
Moreover, non-trivial environment partitions or labels are required for performance guarantee of these OOD objectives (Arjovsky et al., 2019; Krueger et al., 2021; Sagawa* et al., 2020; Ahuja et al., 2021).
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 (Morris et al., 2020; Dwivedi et al., 2020;
Hu et al., 2020).
Alternative options such as random partitions tend not to alleviate the issue (Creager et al., 2021;
Lin et al., 2022), as it can be trivially deemed as mini-batching.
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 (Xu et al., 2018; 2019; Yang et al., 2021).
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 (Xu et al., 2018;
Li et al., 2018a; Xu et al., 2019;
Yang et al., 2021).
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 difﬁcult 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 (Xu et al., 2019), distribution shifts introduced through Z sA and Z sX will doubly mix at the learned features.
In the worst case, the information about Z c A and Z c X could be partially covered by or even replaced by Z s A and Z s X .
This will make OOD generalization of message passing GNNs trained through ERM much more difﬁcult or even impossible.
Besides, as the node representations of 1 ≤ i ≤ k -th layer can also encode graph topology features (Xu et al., 2019), 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 difﬁculty 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 (Xu et al., 2019).
After that, many follow-up variants are proposed to improve the expressivity of GNNs (Morris et al., 2021).
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, there exists many evidences showing that neural networks can fail to generalize to OOD data without a proper objective (Beery et al., 2018;
DeGrave et al., 2021;
Arjovsky et al., 2019; Sagawa* et al., 2020; Bengio et al., 2020; Krueger et al., 2021; Creager et al., 2021; Koyama & Yamaguchi, 2020; Ahuja et al., 2021).
Empirically, we use k -GNNs (Morris et al., 2019) to verify the intuition and observe similar failures for this provably more expressive GNN as basic GNN variants.
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
c
[Figure]
To explore the behaviors of aforementioned methods against complicated distribution shifts on graphs, we ﬁrst modify construction method in Wu et al. (2022c) to construct dataset for Fig.
6, where only FIIF structure-level spurious correlations are injected.
Then we also inject FIIF attribute-level shifts, by setting the node attributes to constant vectors which is spuriously correlated with the labels.
Furthermore, in Fig.
7, graph size shifts are added, which is exactly the SPMotif datasets used in DIR (Wu et al., 2022c).
Besides, in Fig. 8, we can also change the FIIF attribute-level shifts to PIIF attribute-level shifts, where we ﬂip the labels by a probability of 5% and let the ﬂipped label to be spuriously correlated with the node features, following the PIIF SCM in Fig. 4.
Graph size shifts can also be injected in this case, shown as Fig.
Next, we summarize our ﬁndings from the experiments.
Observation I: All existing methods are sensitive to distribution shifts.
From the Fig.
6, 7, 8, 9, we can observe that all
GNNs are sensitive to distribution shifts.
As the intensity of spurious correlation grows, GNNs are more likely to overﬁt to shortcuts presented either in the structure-level or attribute-level, which is similar to general deep learning models (Geirhos et al., 2020).
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.
6, 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 difﬁcult, 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.
4, distribution shifts can happen at both structure-level and attribute-level, and each of them can have different type of spurious correlation with the label.
In Fig.
6, we can ﬁnd 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 veriﬁes the observations in Knyazev et al.
(2019) 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.
6 to Fig.
8, as well as Fig.
7 to Fig.
Observation IV: Using more powerful architectures can not improve the OOD performance.
From the sub-ﬁgures (b) and (c) in Fig.
6, 7, 8, 9, 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.
Speciﬁcally, adding more message passing runs show certain robustness against distribution shifts since they are more likely to learn the desired information during the optimization (Xu et al., 2021a).
However, when the intensity of spurious correlation grows stronger, deeper GNNs are more likely to overﬁt 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.
6, 7, 8, 9 we use 1 - 2 - 3 -GNN following the algorithm of k -GNNs which is provably more expressive than 2 -WL test (Morris et al., 2019).
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 reﬂects 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. 2.
A motivating example.
To begin with, we follow Ahuja et al.
(2021) to introduce a formal example on the failures of GNNs optimized with ERM or IRM (Vapnik, 1991;
Arjovsky et al., 2019) via a linear binary classiﬁcation problem:
Deﬁnition D.1 (Linear classiﬁcation structural equation model (FIIF)) .
[Equation]
where w ∗ inv ∈ R n c with (cid:107) w ∗ inv (cid:107)
= 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 A.1, and latent space interaction as Assumption A.2 or A.3, and strictly separable invariant features A.5, consider a k -layer linearized GNN ρ ◦ h using mean as READOUT for binary graph classiﬁcation, if ∪ e ∈E te supp ( P e ) (cid:54)⊆
[Equation]
: (i) For graphs features generated as Deﬁnition D.1, ρ ◦ h optimized with ERM or IRM will fail to generalize OOD (Eq. 11) almost surely; (ii) For graphs with more than two nodes, globally same node features generated as Deﬁnition 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. 11) almost surely;
For graph classiﬁcation, if the number of nodes is ﬁxed to one, it covers the linear classiﬁcation as above.
When ∪ e ∈E te supp ( P e ) (cid:54)⊆ ∪ 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
[Equation]
[Equation]
) > 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., ﬁnding the ground truth w ∗ inv , following the Theorem 3 from Ahuja et al.
Theorem D.2 (Insufﬁciency of ERM and IRM) .
Suppose each e ∈ E all follows Deﬁnition.
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 inﬁnite 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 (ﬁnding 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 Deﬁnition 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 (Sagawa* et al., 2020; Krueger et al., 2021; Ahuja et al., 2021; Bellot & van der Schaar, 2020) 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.1 of Rosenfeld et al.
(2021), when there is not sufﬁcient 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 (Creager et al., 2021) 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 ﬁnd a subset of the inputs for interpretable prediction robustly against distribution shifts.
Here we also provide a discussion for these works.
The ﬁrst work following this line is I NV RAT (Chang et al., 2020), 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 (Wu et al., 2022c), 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 justiﬁcation established for DIR (Theorem 1 to Corollary 1 in Wu et al. (2022c)) 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 ﬁnd the unstable and relatively high sensitivity of DIR to spurious correlations, which veriﬁes our ﬁnding.
More details about empirical behaviors of DIR can be found in Appendix G.
In contrast to DIR, GIB (Yu et al., 2021) 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 (Wu et al., 2022c) as they also focus on FIIF), which further implies Y ⊥⊥ S | ˆ G c .
Hence it is natural to use Information Bottleneck (IB) objective (Tishby et al., 1999) to solve for G c :
[Equation] [Equation]
which explains the success of many existing works in ﬁnding predictive subgraph through IB (Yu et al., 2021).
However, the estimation of I ( ˆ G c , G ) is notoriously difﬁcult due to the complexity of graph, which can lead to unstable convergence as observed in our experiments.
In contrast, optimization with contrastive objective in GOOD as Eq.
8 induces more stable convergence.
E. Theory and Discussions
In this section, we provide proofs for propositions and theorems mentioned in the main paper.
E.1. Challenges of OOD generalization on graphs.
From the aforementioned analysis, we can summarize some key challenges revealed by the 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.2. Proof for theorem B.2 (i)
Theorem E.1 (GOODv1 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.1, assuming that (a) f G gen and f G c gen in Assumption A.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.
7, is an invariant GNN (Def. B.1).
Proof.
We re-write the objective as follows:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
The proof of Theorem E.1 is essentially to show the estimated ˆ G c through Eq.
17 is the underlying G c , then the maximizer of I ( ˆ G c ; Y ) in Eq.
17 can produce most informative and stable predictions about Y based on G , hence is an invariant GNN according to Deﬁnition.
B.1.
In the next, we are going to take a information-theoretic view of the ﬁrst term
I ( ˆ G c ; Y ) and the second term I ( ˆ G c ; ˜ G c | Y ) to conclude the proof.
We begin by introducing the following lemma:
Lemma E.2.
Given the same conditions as Thm.
E.1, I ( ˆ G c ; Y ) is maximized if and only if I ( ˆ G c ; Y | E = e ) is maximized, ∀ e ∈ E tr .
The proof for Lemma E.2 is straightforward, given the condition that samples from each training environment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr .
Obviously, ˆ G c =
G c is a maximizer of I ( ˆ G c ; Y )
=
I ( C ; Y ) =
H ( Y ) , since f c gen : C → G c is invertible and C causes Y .
However, there might be some subset G ps ⊆
G s from the underlying G s that entail the same information about label, i.e., I ( G pc ∪ G ps ; Y )
=
I ( G c ; Y ) where ˆ G c = G pc ∪ G ps and G pc = G c ∩ ˆ G c .
For
FIIF (Assumption 4(b)), it can not happen otherwise, let G lc
= G c − G pc , then we have:
[Equation]
where the second last equality is due to C → Y and the invertibility of f c gen : C → G c in FIIF, i.e., H ( Y | C )
=
H ( Y | G c ) =
H ( Y | G pc , G lc ) = 0 .
However, in PIIF, it can not hold since conditioning on G pc , G ps can not determine Y , since S (cid:54)⊥⊥ Y | C .
In other words, G s (cid:54)⊥⊥ 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 ˆ G c , we will use the second term to eliminate it:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ 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 ( ˆ G c ; ˜ G c | Y ) is essentially maximizes
I ( ˆ G c , E = ˆ e ; ˜ G c , E = ˜
e | Y ) , ∀ ˆ e, ˜ e ∈ E tr .
Hence, we have:
[Equation]
[Equation]
We claim
Eq.
can eliminate any potential subsets in the estimated ˆ G c .
[Equation]
Otherwise, suppose there are some subsets ˆ G ps ⊆
ˆ G s and ˜ G ps ⊆
˜ G s contained in the estimated ˆ G c , ˜ G c , where ˆ
G s , ˜ G s be the corresponding underlying G s s for ˆ G c , ˜ G c .
Let ˆ G ∗ c and ˜
G ∗ c be the ground truth invariant subgraph G c s of ˆ G and ˜ G , ˆ G lc
= ˆ G ∗ c − ˆ G c and ˜ G lc =
˜
G ∗ c − ˜ G
c be the l eft (un-estimated) subsets from corresponding ground truth G c s, and
ˆ G pc
=
ˆ G ∗ c − ˆ G lc and ˜ G pc = ˜
G ∗ c − ˜
G lc be the complement, or equivalently, the p artial ˆ G ∗ c , ˜
G ∗ c that are estimated in ˆ G c , ˜ G c , respectively.
We can also deﬁne similar counterparts for G s : ˆ G ps , ˜ G ps are the partial ˆ G s , ˜ G s s contained in the estimated ˆ G c , ˜ G c while ˆ G ls , ˜ G ls are the left subsets ˆ G s , ˜ G s , respectively.
Recall the constraint that | G c | =
s c , hence if ˆ G pc ⊆
ˆ G c , then a corresponding ˆ G lc
= ˆ G ∗ c − ˆ G pc will be replaced by ˆ G ps in ˆ G c .
In this case, we have:
[Equation]
where the second equality is due to E = ˆ e is determined so that H ( E = ˆ e
| ˆ G c , Y ) = 0 .
Compared Eq.
21 to that when ˆ
G c = ˆ G ∗ c , we have the entropy change as:
[Equation]
Let (cid:15) =
H ( ˆ G ps | ˆ G pc , E = ˆ e, Y ) .
In a idealistic setting, when the noise of the generation process S := f spu ( Y, E ) in PIIF tends to be 0 , i.e., (cid:15) → 0 , S is determined conditioned on E, Y , hence G s and any subsets of G s are all determined.
Then, it sufﬁces to know that in Eq.
22, H ( ˆ G ps | ˆ G pc , E =
ˆ e, Y )
=
while H ( ˆ G lc | ˆ G pc , E = ˆ e, Y ) > 0 since ˆ G lc can not be determined when given ˆ G pc , E = ˆ e, Y .
Thus, when some subset from G s is included in ˆ G c , it will minimize H ( ˆ G c , E
= ˆ e | Y ) .
In the next, we will show how (cid:15) =
H ( ˆ G ps | ˆ G pc , E = ˆ e, Y ) can be cancelled thus leading to a smaller H ( ˆ G c , E
= ˆ e | Y ) , by considering the second term H ( ˆ G c , E
= ˆ e | ˜ G c , E = ˜
e, Y ) .
As for H ( ˆ G c , E = ˆ e | ˜ G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possible cases into two:
(i) One of ˆ G c and ˜ G c contains some subset of G s , i e., ˆ G c contains some ˆ G ps ⊆
ˆ G s ; (ii)
Both ˆ G c and ˜ G c contain some ˆ G ps ⊆
ˆ G s and ˜ G ps ⊆
˜ G s , respectively.
For (i), we have:
[Equation]
Thus, H ( ˜ G ps | ˆ G c , E = ˜
e, Y, ˜
G lc , E = ˆ e ) =
and Thus, we have:
[Equation]
Combing ∆ H ( ˆ G c , E
= ˆ e | Y ) , we have:
[Equation]
where the last equality is because of the independence of ˆ G ps between ˜ G c , E = ˜
e conditioned on Y, E = ˆ e .
Since conditioning will lower the entropy for both discrete and continuous variables (Cover & Thomas, 2006; Yeung, 2008), we have:
[Equation]
which implies the existence of ˆ G ps in ˆ G c will lower down the second term in Eq.
17 for the case (i)
[Equation]
For (ii), we have:
[Equation]
Similar to (i), H ( ˆ G ps | ˜ G pc , ˜ G ps , E = ˜
e, Y, ˆ G pc , E = ˆ e ) can be cancelled out with H ( ˆ G ps | ˆ G pc , E = ˆ e, Y ) .
Then, we have:
[Equation]
Since additionally conditioning on ˆ G ps in H ( ˆ G lc , E = ˆ e | ˜ G pc , ˜ G ps , E = ˜
e, Y ) can not lead to new information about ˆ G lc , we have:
[Equation]
which follows that ∆ I ( ˆ G c , E = ˆ e ; ˜ G c , E = ˜
e | Y ) < 0 .
To summarize, the ground truth G c is the only maximizer of the objective (Eq. 17), hence solving for the objective (Eq. 17) can elicit an invariant GNN.
E.3. Proof for theorem B.2 (ii)
Theorem E.3 (GOODv2 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.1, assuming that (a) f G gen and f G c gen in Assumption A.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.
7, is an invariant GNN (Def. B.1).
Proof.
We re-write the objective as follows:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
Similar to the proof for Theorem E.1, to prove Theorem E.3 is essentially to show the estimated ˆ G c through Eq.
30 is the underlying G c , hence the minimizer of Eq.
30 elicits an invariant GNN predictor.
In the next, we also begin with a lemma:
Lemma E.4.
Given data generation process as Theorem E.3, for both FIIF and PIIF, we have:
I ( C ; Y ) ≥ I ( S ; Y ) , hence I ( G c ; Y ) ≥ I ( G s ; Y ) .
Proof for Lemma E.4.
For both FIIF and PIIF, Assumption A.5 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.4, we know ˆ G c at least contains some subset of the underlying G c , otherwise the constraint I ( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) will be violated since G c ⊆
ˆ G s in this case.
Assuming there are some subset of G s contained in ˆ G c , without loss of generality, we can divide all of the possible cases about ˆ G c into two:
(i) ˆ G c only contains a subset of the underlying G c ; (ii) ˆ G c contains a subset of the underlying G c as well as part of the underlying G s ;
[Equation]
Before the discussion, let us inherit the notations of subsets of G c , G s from the proof for Theorem E.1: Let ˆ G ∗ c and ˜
G ∗ c be the ground truth invariant subgraph G c s of ˆ G and ˜ G , ˆ G lc
= ˆ G ∗ c − ˆ G c and ˜ G lc =
˜
G ∗ c − ˜ G
c be the l eft (un-estimated) subsets from corresponding ground truth G c s, and
ˆ G pc
=
ˆ G ∗ c − ˆ G lc and ˜ G pc = ˜
G ∗ c − ˜
G lc be the complement, or equivalently, the p artial ˆ G ∗ c , ˜
G ∗ c that are estimated in ˆ G c , ˜ G c , respectively.
Similarly, ˆ G ps , ˜
G ps are the partial ˆ G s , ˜ G s s contained in the estimated ˆ G c , ˜ G c while ˆ G ls , ˜ G ls are the left subsets ˆ G s , ˜ G s , respectively.
First of all, case (i) cannot hold because, when maximizing I ( ˆ G c ; ˜ G c | Y ) , if ∃ ˆ G lc
= ˆ G ∗ c − ˆ G c , as shown in the proof for Theorem E.1, including ˆ G lc into ˆ G c can always enlarge I ( ˆ G c ; ˜ G c | Y ) , while not affecting the optimality of I ( ˆ G s ; Y )
+ I ( ˆ G c ; Y ) by re-distributing ˆ G lc from ˆ G s to ˆ G c .
Consequently, ˆ G ∗ c must be included in ˆ G c , i.e., ˆ
G ∗ c ⊆ ˆ G c .
As for case (ii), recall that, by the condition of equally distributed training samples from each training environment, maximizing I ( ˆ G c ; ˜ G c | Y ) is essentially maximizing I
( ˆ G c , E = ˆ e ; ˜ G c , E =
˜
e | Y ) , ∀ ˆ e, ˜ e ∈ E tr , hence, we have:
[Equation]
We claim Eq.
31 can eliminate any potential subsets in the estimated ˆ G c .
Similarly, we have:
[Equation]
where the second equality is due to E
= ˆ e is determined.
Compared to the case that ˆ G c
= ˆ G ∗ c , we have:
[Equation]
Then, as for H ( ˆ G c , E = ˆ e | ˜ G c , E = ˜
e, Y ) , without loss of generality, we can divide all of the possible cases into two:
(a) ˆ G c contains some ˆ G ps ⊆
ˆ G s ; (b) Both ˆ G c and ˜ G
c contain
some
ˆ G ps ⊆ ˆ G s and ˜ G ps ⊆
˜ G s , respectively.
For (a), we have:
[Equation]
Similarly to the proof for Theorem E.1, when considering ∆ I ( ˆ G c ; ˜ G c | Y ) , the effects of H ( ˆ G ps | ˜ G c , E = ˜
e, Y, ˆ G ∗ c , E = ˆ e ) is cancelled out by H ( ˆ G ps | ˆ G ∗ c , E = ˆ e, Y ) .
Hence, we have:
[Equation]
For (b), we have:
[Equation]
Similarly, H ( ˆ G ps | ˜ G ∗ c , ˜ G ps , E = ˜
e, Y, ˆ G ∗ c , E = ˆ e ) = 0 can also be cancelled out by H ( ˆ G ps | ˆ G ∗ c , E = ˆ e, Y ) .
Moreover, for
H ( ˆ G ∗ c | ˜ G ∗ c , ˜ G ps , E = ˜
e, Y, E
= ˆ e ) , ˜ G ps can not bring no additional information about ˆ G ∗ c , when conditioning on ˜ G ∗ c , Y, E = ˜
e .
Hence, we also have:
[Equation]
To summarize, when maximizing I ( ˆ G c ; ˜ G c | Y ) , including any ˆ G ps ⊆
ˆ G ∗ s can not bring additional beneﬁt while affecting the optimality of I ( ˆ G s ; Y )
+ I ( ˆ G c ; Y ) .
More speciﬁcally, when considering the changes to I ( ˆ G s ; Y )
+ I ( ˆ G c ; Y ) ,
∀ G ps ⊆
G s , we have
[Equation]
while I ( Y ; ˆ G ∗ c , G p s )
=
I ( Y ; ˆ G ∗ c ) + I
( Y ; ˆ G ps | ˆ G ∗ c ) , ∀ e ∈ E tr .
Consequently,
[Equation]
Hence, only the underlying G c is the solution to Eq. 30, which implies that solving for the objective (Eq. 30) can elicit an invariant GNN.
F. Details of Prototypical GOOD Implementation
In fact, the GOOD framework introduced in Sec.
3 can have multiple implementations.
We choose interpretable architectures in our experiments for the purpose of concept veriﬁcation.
More sophisticated architectures can be incorporated.
Experimental results in Sec. 4 also demonstrates that, even equipped with basic GNN architectures, GOOD 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 GOOD decomposes a GNN model for graph classiﬁcation into two modules, i.e., a featurizer:
g : G → G c and a classiﬁer f c : G c → Y .
Speciﬁcally, for the implementation of Featurizer, we choose one of the common practices GAE (Kipf & Welling, 2016) for calculating the sampled weights for each edge.
More formally, the soft mask is predicted through the following equation:
[Equation]
If a sampling ratio s c is predetermined, we sample s c of total edges with the largest predicted weights as a soft estimation of ˆ
G c .
Then, the estimated ˆ G c will be forwarded to the classiﬁer f c for predicting the labels of the original graph.
Although
Theorem E.1 assumes s c is known, in real applications we do not know the speciﬁc s c .
Hence, in experiments, we select s c according to the validation performance.
To thoroughly study the effects of I ( ˆ G s ; Y ) comparing to GOODv1, we stick to using the same s c and sampling process for GOODv2, while GOODv2 essentially requires less speciﬁc knowledge about ground truth r c hence achieving better empirical performance.
Moreover, once the sampled edges are determined, the classiﬁer GNN can take either the original feature of the input graph or the learned feature from the featurizer as the new node attributes for ˆ 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 GOODv2 for elaboration while the implementation of GOODv1 can be obtained via removing the third term from GOODv2.
Recall that GOODv2 has the following formulation:
[Equation]
where ˆ G c = g ( G ) , ˜ G c = g ( ˜ G ) and ˜ G ∼ P ( G | Y ) , i.e., ˜ G and G have the same label.
In Sec.
B.4, we introduce a contrastive approximation for I ( ˆ G c ; ˜ G c | Y ) : I
[Equation]
where positive samples ( ˆ G c , ˜ G c ) are the extracted subgraphs of graphs that have the same label of G , negative samples are those with different labels, P g ( G |Y
=
Y ) is the pushforward distribution of P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label
Y , h ˆ G c , h ˜ G c , h G ic are the graph presentations of the estimated subgraphs, and φ is the similarity metric for the graph presentations.
As M → ∞ , Eq.
38 approximates I ( ˆ G c ; ˜ G c | Y ) which can be regarded as a non-parameteric resubstitution entropy estimator via the von Mises-Fisher kernel density (Ahmad & Lin, 1976;
Kandasamy et al., 2015; Wang & Isola, 2020).
While for the third term I ( ˆ G s ; Y ) and the constraint
I
( ˆ G s ; Y ) ≤
I ( ˆ G c ; Y ) , a straightforward implementation is to imitate the hinge loss:
[Equation]
[Equation]
where N is the number of samples, I is a indicator function that outputs 1 when the interior condition is satisﬁed otherwise 0 , and R ˆ G s and R ˆ G c are the empirical risk vector of the predictions for each sample based on ˆ G s and ˆ G c respectively.
One can also formulate Eq.
37 from game-theoretic perspective (Chang et al., 2020).
Finally, we can derive the speciﬁc loss for the optimization of GOODv2 combining Eq.
38 and Eq. 39:
[Equation]
where R ˆ G c , R ˆ G s are the empirical risk when using ˆ G c , ˆ G s to predict Y through the classiﬁer.
Typically, we use a additional
MLP downstream classiﬁer ρ s for ˆ G s in the classiﬁer GNN.
h ˆ G c is the graph representation of ˆ G c which can be induced from the GNN encoder either in the featurizer or in the classiﬁer.
α, β are the weights for I ( ˆ G c ; ˜ G c | Y ) and I ( ˆ G s ; Y ) , and φ is implemented as cosine similarity.
The optimization loss for GOODv1 merely contains the ﬁrst two terms in Eq.
The detailed algorithm for GOOD is given in the Algorithm 1, assuming the h ˆ G c is obtained via the graph encoder in f c .
G. Full Experimental Results and Settings
In this section, we provide more details about the experimental results and settings in complementary to Sec. 4, including the full empirical results and analysis, dataset preparation, dataset statistics, implementations of baselines, selection of models and hyperparameters as well as evaluation protocols.
G.1. Full experimental results and analysis
Datasets.
We use the SPMotif datasets from DIR (Wu et al., 2022c) where artiﬁcial structural shifts and graph size shifts are nested (SPMotif-Struc).
Besides, we construct a harder version mixed with attribute shifts (SPMotif-Mixed).
To examine
GOOD in real-world scenarios with more complicated relationships and distribution shifts, we also use DrugOOD (Ji et al., 2022) from AI-aided Drug Discovery with Assay, Scaffold, and Size splits, convert the ColoredMNIST from IRM (Arjovsky et al., 2019) using the algorithm from Knyazev et al.
(2019) to inject attribute shifts, and split Graph-SST (Yuan et al., 2020) to inject degree biases.
To compare with previous specialized OOD methods for graph size shifts (Yehudai et al., 2021;
Bevilacqua et al., 2021), we use the datasets in Bevilacqua et al.
(2021) that are converted from TU benchmarks (Morris et al., 2020).
Baselines and our methods.
Besides the ERM, we also compare with SOTA interpretable GNNs, GIB (Yu et al., 2021),
ASAP Pooling (Ranjan et al., 2020), and DIR (Wu et al., 2022c), to validate the effectiveness of the optimization objective in GOOD.
We use the same selection ratio (i.e., s c ) for all models.
On the other hand, to validate the effectiveness of the decomposition in GOOD, we compare GOOD with SOTA OOD objectives including IRM (Arjovsky et al., 2019), v-Rex (Krueger et al., 2021) and IB-IRM (Ahuja et al., 2021), for which we apply random environment partitions.
We also compare GOOD with EIIL (Creager et al., 2021) and CNC (Zhang et al., 2022) that does not require environment labels, where CNC (Zhang et al., 2022) has a more sophisticated contrastive sampling strategy for combating subpopulation shifts.
OOD performance on structure and mixed shifts.
In Table 3, 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 converge 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.2 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, GOODv1 and GOODv2 outperform all of the baselines by a signiﬁcant margin up to 10% with lower variance , which demonstrates the effectiveness and excellent OOD generalization ability of GOOD.
More analysis and results on real-world datasets are given in Appendix G.
[Table] [Table]
Lower rank is better.
OOD generalization performance on realistic shifts.
In Table 4 and Table 5, we examine the effectiveness of GOOD 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 (Gulrajani & Lopez-Paz, 2021; Ji et al., 2022).
However, both GOODv1 and GOODv2 consistently and signiﬁcantly outperform previous methods, including previous specialized methods Γ GNNs (Bevilacqua et al., 2021) for combating graph size shifts, demonstrating the generality and superiority of GOOD.
OOD generalization performance on graph size shifts.
In Table 5, we additionally compare GOOD to the specialized designed methods for OOD generalization in terms of graph sizes ( Γ GNNs)
(Bevilacqua et al., 2021), where we include the author reported results for both kernel methods and Γ GNNs.
It can be found that GOOD consistently and signiﬁcantly outperforms the previous SOTA methods, which further demonstrates the generality of GOOD.
[Table]
Comparisons with advanced ablation variants.
As discussed in Sec.
B.4, GOOD 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 GOOD.
Furthermore, although GOODv1 can outperform
GOODv2 when we may have a relatively accurate s c , the improvements in GOODv1 are not as stable as GOODv2 or even unsatisfactory when the assumption is violated.
This phenomenon also reveals the advances of GOODv2 in practical scenarios.
G.2. 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 6.
SPMotif datasets.
We construct 3-class synthetic datasets based on BAMotif (Ying et al., 2019; Luo et al., 2020) following (Wu et al., 2022c), 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, 1000 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 on FIIF, where the motif and one of the three base graphs (Tree, Ladder, Wheel) are artiﬁcially (spuriously) correlated with a probability of various biases, and equally correlated with the other two.
Speciﬁcally, given a predeﬁned bias b , the probability of a speciﬁc motif (e.g., House) and a speciﬁc 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 inﬂuences 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.
Speciﬁcally, given a predeﬁned 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.
(2021); Bevilacqua et al.
(2021) to study the
OOD generalization abilities of various methods on four of TU datasets (Morris et al., 2020), i.e., PROTEINS, DD, NCI1,
NCI109 .
Speciﬁcally, we use the data splits generated by Yehudai et al.
and use the Matthews correlation coefﬁcient as evaluation metric following (Bevilacqua et al., 2021) 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 7.
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 (Yuan et al., 2020), that converts sentiment sentence classiﬁcation datasets SST5 and
SST-Twitter (Socher et al., 2013; Dong et al., 2014) into graphs, where node features are generated using BERT (Devlin et al., 2019) and the edges are parsed by a Biafﬁne parser (Gardner et al., 2018).
Our splits are created according to the averaged degrees of each graph.
Speciﬁcally, 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 (Arjovsky et al., 2019).
We convert the ColoredMnist into graphs using super pixel algorithm introduced by Knyazev et al.
Speciﬁcally, 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 ﬂipped 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 ﬂipped 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 afﬁnity prediction for both macromolecule (protein target) and small-molecule (drug compound).
The molecule data and the notations are curated from realistic ChEMBL database (Mendez et al., 2019).
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 Afﬁnity Prediction which uses ic50 measurement type and contains core level annotation noises.
For more details, we refer interested readers to Ji et al.
[Table]
[Table]
G.3. 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 (Ioffe & Szegedy, 2015) between layers and JK residual connections at last layer (Xu et al., 2018).
For the architectures we use the GCN with mean readout (Kipf & Welling, 2017) for all datasets except Proteins where we empirically observe better validation performance with a GIN and max readout (Xu et al., 2019), and for DrugOOD datasets where we follow the backbone used in the paper (Ji et al., 2022), i.e., 4 -layer GIN with sum readout.
The hidden dimensions are ﬁxed 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 (Kingma & Ba, 2015) with a learning rate of 1 e −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 (Ji et al., 2022).
To avoid underﬁtting, we pretrain models for 20 epochs for all datasets, except for CMNIST and Twitter where we pretrain 5 epochs and for SST5 we pretrain 10 epochs, because of the dataset size and the difﬁculty of the task.
To avoid overﬁtting, we also employ an early stopping of 5 epochs according to the validation performance.
Meanwhile, dropout (Srivastava et al., 2014) is also adopted for some datasets.
Speciﬁcally, 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 (Yu et al., 2021; Ranjan et al., 2020), where we use the codes provided by the authors 3 for DIR c(Wu et al., 2022c) 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.
(2022c), 0 .
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 ﬁnal outcomes.
Although SPMotif-Struc is also evaluated in DIR, we ﬁnd 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 line 278 in 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 (Gulrajani & Lopez-Paz, 2021) for IRM (Arjovsky et al., 2019), V-Rex (Krueger et al., 2021) and IB-IRM (Ahuja et al., 2021).
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 (Creager et al., 2021), 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 (Zhang et al., 2022), we follow the algorithm description to modify the sampling strategy in supervised contrastive loss (Khosla et al., 2020) based on a pretrained GNN optimized with ERM, and choose the weight for contrastive loss using the same grid search as for GOOD.
Implementations of GOOD.
For fair comparison, GOOD 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 ﬁx 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 ﬁnd 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 ﬁnd that increasing α does not help increase the performance with few random runs.
Besides, we also have various implementation options for obtaining the features in ˆ G c , for obtaining h ˆ G c , as well as for obtaining predictions based on ˆ G s .
By default, we feed the graph representations of featurizer GNN to the classiﬁer GNN, as well as to the contrastive loss.
For classifying G based on ˆ G s , we use a separate MLP downstream classiﬁer in the classiﬁer 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 overﬁt to the shortcuts hence we instead feed the original features to the downstream classiﬁer GNN.
There can be more other options, such as using separate graph convolutions on ˆ
G s or ˆ G c , which we leave for future work.
Options for obtaining the features in ˆ G c are: { from g , from the raw features } .
Options for obtaining h ˆ G c are: { from the GNN encoder of the classiﬁer f c with the same pooling as the classiﬁer, from the
GNN encoder of the featurizer g with a SUM global pooling, } .
Options for obtaining predictions based on ˆ G s are: { from another classiﬁer with shared GNN encoder of f c , from another classiﬁer with shared GNN encoder of f c while without gradients backwards to the encoder, from a single GNN convolution and a same pooling as f c } .
We select the corresponding options according to the validation performance with several runs of random α and β , and stick to one for each dataset.
As a result, we empirically ﬁnd using the raw node features for ˆ G c , obtaining h ˆ G c via a global ADD pooling with the featurizer outputs, and obtaining predictions based on ˆ G s from another classiﬁer with shared GNN encoder of f c while without gradients backwards to the encoder, has better validation performances.
Except for TU datasets where we use the outputs of g as the features of ˆ G c , and obtain predictions with one GNN layer for the prediction of ˆ G s empirically has better validation performances.
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.4. Software and hardware
We implement our methods with PyTorch (Paszke et al., 2019) and PyTorch Geometric (Fey & Lenssen, 2019).
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.