Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs
Abstract
Despite recent success in using the invariance principle for out-of-distribution (OOD) generalization on Euclidean data (e.g., images), studies on graph data are still limited.
Different from images, the complex nature of graphs poses unique challenges to adopting the invariance principle.
In particular, distribution shifts on graphs can appear in a variety of forms such as attributes and structures, making it difficult to identify the invariance.
Moreover, domain or environment partitions, which are often required by OOD methods on Euclidean data, could be highly expensive to obtain for graphs.
To bridge this gap, we propose a new framework, called C ausality
Inspired I nvariant G raph Le A rning (CIGA), to capture the invariance of graphs for guaranteed OOD generalization under various distribution shifts.
Specifically, we characterize potential distribution shifts on graphs with causal models, concluding that OOD generalization on graphs is achievable when models focus only on subgraphs containing the most information about the causes of labels.
Accordingly, we propose an information-theoretic objective to extract the desired subgraphs that maximally preserve the invariant intra-class information.
Learning with these subgraphs is immune to distribution shifts.
Extensive experiments onsynthetic or real-world datasets, including a challenging setting – DrugOOD,from
AI-aided drug discovery, validate the superior OOD performance of CIGA 1 .
Introduction
Graph representation learning with graph neural networks (GNNs) has gained great success in tasks involving relational information [45, 35, 99, 106, 107].
However, it assumes that the training and test graphs are drawn from the same distribution, which is often violated in reality [37, 47, 38, 40].
The mismatch between training and test distributions, i.e., distribution shifts , introduced by some underlying environmental factors related to data collection or processing, could seriously degrade the performance of deployed models [7, 24].
Such out-of-distribution (OOD) generalization failures become the major roadblock for practical applications of graph representation learning [40].
Meanwhile, enabling OOD generalization on regular Euclidean data has received surging attention and several solutions were proposed [4, 81, 10, 49, 23, 48, 2].
In particular, the invariance principle from causality is at the heart of those works [76, 74, 79].
The principle leverages the Independent
Causal Mechanism (ICM) assumption [74, 77] and implies that, model predictions that only focus on the causes of the label can stay invariant to a large class of distribution shifts [76, 4, 2].
“Cycle” “House”
Classification Invariant Subgraph Identification
[Figure]
Algorithm
OODGuarantee Regime E Known SCMSupport
IRM[4]
Yes
R
Yes
PIIF IB-IRM[2]
Yes
R
Yes
PIIF&FIIF EIIL[23]
Yes
R
No
PIIF DANN[31]
N/A
R
Yes
N/A MatchDG[61]
N/A
R
Yes
FIIF GroupDro[81]
N/A
R
Yes
N/A CNC[124]
N/A
R
No
N/A GIB[120]
Yes
G
No
FIIF DIR[104]
No
G
No
FIIF CIGA (Ours)
Yes
G
No
PIIF&FIIF
[Figure]
Despite the success of the invariance principle on Euclidean data, the complex nature of graphs raises several new challenges that prohibit direct adoptions of the principle.
First, distribution shifts on graphs are more complicated.
They can happen at both attribute-level and structure-level, and be observed in multiple forms such as graph sizes, subgraph densities and homophily [113, 11, 102].
On the other hand, each of the shifts can spuriously correlate with labels in different modes [4, 71, 2].
Consequently, the entangled complex distribution shifts make it more difficult to identify and capture the invariance on graphs.
Second, OOD algorithms developed and analyzed on Euclidean data often require additional environment (or domain) labels for distinguishing the sources of distribution shifts [4].
However, the environment labels could be highly expensive to obtain and thus often unavailable for graphs, as collecting the labels usually requires expert knowledge due to the abstraction of graphs [37].
These challenges render the problem studied in this paper even more challenging:
How could one generalize the invariance principle to enable OOD generalization on graphs?
To solve the above problem, we propose C ausality
Inspired I nvariant G raph Le A rning (CIGA), a new framework for capturing the invariance of graphs to enable guaranteed OOD generalization under different distribution shifts.
Specifically, we build three Structural Causal Models (SCMs) [74] to characterize the distribution shifts that could happen on graphs: one is to model the graph generation process, and the other two are to model two possible interactions between invariant and spurious features during the graph generation, i.e., Fully Informative Invariant Feature (FIIF) and
Partially
Informative Invariant Feature (PIIF) (Sec. 2.2).
Then, we generalize the invariance principle to graphs for OOD generalization: GNN models are invariant to distribution shifts if they focus only on an invariant and critical subgraph G c that contains the most of the information in G about the underlying causes of the label.
Thus, the problem of achieving OOD generalization on graphs can be rephrased into two processes: invariant subgraph identification and label prediction.
Accordingly, shown as Fig. 1(a), we introduce a prototypical invariant graph learning algorithm that decomposes a GNN into: a) a featurizer g for identifying the underlying invariant subgraph G c from G ; b) a classifier f c for making predictions based on G c .
To extract the desired subgraph G c , we derive an information-theoretic objective for the featurizer to identify subgraphs that maximally preserves the invariant intra-class information across a set of different (unknown) environments.
We theoretically show that this approach can provably identify the underlying G c under mild assumptions (Sec. 3).
Experiments on 16 synthetic and real-world datasets with various distribution shifts, including a challenging setting from AI-aided drug discovery [40], show that CIGA can significantly outperform all of existing methods up to 10% , demonstrating its promising OOD generalization ability (Sec. 4).
Related Work.
We review existing methods that might improve the OOD generalization on graphs, summarize the main differences between our solution and them in Table 1(b), and leave thorough discussions to Appendix B.2.
On Euclidean data, Invariant Learning [4, 23, 2], Group Distributionally
Robust Optimization [49, 81, 124], Domain Adaption and Domain Generalization [31, 93, 52, 27, 61,
100] are three widely adopted approaches to enable OOD generalization.
However, they all have their own limitations when being applied to graphs.
First, previous invariant learning methods are mostly developed and analyzed for Euclidean data [4, 2, 23], or under specific SCM assumptions [4], making the theoretical results hardly able to generalize to the complicated graph data [80] that can have multiple types of distribution shifts [71].
Group Distributionally Robust Optimization that minimizes the gap between worst group risk and average risk [49, 81, 124], and Domain Adaption/Generalization methods that aim to learn class-conditional domain invariant representations [31, 93, 52, 27, 100], cannot guarantee a min-max optimal predictor without additional assumptions [126, 4, 2].
Moreover, most existing methods require environment labels that are however expensive to obtain in graphs, which limits their applications to graphs [4, 49, 2, 81, 31, 93, 27, 61].
In contrast, we aim to develop
OOD algorithms for graphs that are provably generalizable under different types of distribution shifts.
Another line of relevant works is about GNN explainability that aims to find a subgraph of the input as the explanation for a GNN prediction [116, 122].
Although some may leverage causality to justify the generated explanation [53], they mostly focus on understanding the predictions of GNNs instead of for OOD generalization.
The closest works to ours are two interpretable GNNs that aim to explicitly extract a subgraph for both predictions and explanations guided by information theory [120] and causality [104], respectively.
However, they focus on graphs and shifts generated under a specific
SCM.
Although one of them can provide theoretical guarantee for OOD generalization [120] by using the information bottleneck criteria [2], they would inevitably fail to generalize to graphs generated under different SCMs.
More discussions about the failure are deferred to Appendix D.4.
Besides,
Bevilacqua et al.
[11] also discuss OOD generalization on graphs but limited to a specific graph family and graph size shifts.
Wu et al.
[103] propose OOD generalization algorithms on graphs for the task of node classification, also limited to graphs and shifts under a specific SCM.
To the best of our knowledge, there is no existing work that could handle more comprehensive graph distribution shifts than CIGA, while also achieving provable OOD generalization performance.
OOD Generalization on Graphs through the Lens of Causality 2.1 Problem Setup
In this work, we focus on OOD generalization in graph classification.
Specifically, we are given a set of graph datasets D =
{D e } e collected from multiple environments E all .
Samples ( G ei , Y ei ) ∈
D e from the same environment are considered as drawn independently from an identical distribution P e .
A
GNN ρ ◦ h generically has an encoder h : G → R h that learns a meaningful representation h G for each graph G to help predict the label ˆ Y G
= ρ ( h G ) with a downstream classifier ρ : R h
→ Y .
The goal of
OOD generalization on graphs is to train a GNN ρ ◦ h with data from training environments D tr = {D e } e ∈E tr ⊆E all that generalizes well to all (unseen) environments, i.e., to minimize max e ∈E all R e , where R e is the empirical risk of ρ ◦ h under environment
e [97, 4].
We leave more details about the background of GNN for graph classification and invariant learning in Appendix B.1.
It is known that OOD generalization is impossible without assumptions on the environments E all [74, 2].
Thus, we will first formulate the data generation process with structural causal model and latentvariable model [74, 77, 50], to characterize the distribution shifts that could happen on graphs.
Then, we investigate whether the existing methods are generalizable under these distribution shifts.
2.2 Graph Generation Process
[Figure]
We take a latent-variable model perspective on the graph generation process and assume that the graph is generated through a mapping f gen :
Z → G , where Z ⊆ R n is the latent space and G = ∪ ∞ N =
[Equation]
is the graph space.
Let E denote environments.
Following previous works [50, 2], we partition the latent variable from Z into an invariant part C ∈ C =
R n c and a varying part
S ∈ S =
R n s , s.t., n = n c + n s , according to whether they are affected by E or not.
Similarly in images, C and S can represent content and style while E can refer to the locations where the images are taken [7, 125, 50].
Furthermore, C and S control the generation of the observed graphs (Assumption 2.1) and can have multiple types of interactions at the latent space (Assumptions 2.2, 2.3).
Graph generation model.
We elaborate the SCM for the graph generation process in Assumption 2. and Fig. 2(a), where noises in the structural equations are omitted for simplicity [77].
Assumption 2.1 (Graph Generation Structural Causal Model) .
[Equation]
In Assumption 2.1, f gen is decomposed into f G c gen , f G s gen and f G gen to control the generation of G c , G s , and G , respectively.
Among them, G c inherits the invariant information of C that would not be affected by the interventions (or changes) of E [74, 77].
For example, certain properties of a molecule can usually be described by a sub-molecule, or a functional group, which is invariant across different species or assays [12, 92, 40].
On the contrary, the generation of G s and G will be affected by the environment E through S .
Thus, graphs collected from different environments (or domains) can have different distributions of structure-level properties (e.g., graph sizes [11, 102]) as well as feature-level properties (e.g., homophily [62, 17]).
Therefore, the subgraph G s inherits the spurious feature about
Y [125].
In fact, Assumption 2.1 is compatible with many graph generation models by specifying the function classes of f G c gen , f G s gen and f G gen [89, 57, 117, 59].
Since our goal is to characterize the potential distribution shifts in Assumption 2.1, we focus on building a general SCM that is compatible to many graph families and leave graph family specifications and their implications to OOD generalization in future works.
More discussions are provided in Appendix C.
Interactions at latent space.
Following previous works [4, 2], we categorize the latent interactions between C and S into Fully Informative Invariant Features (FIIF, Fig. 2(b)) and Partially Informative
Invariant Features (PIIF, Fig. 2(c)) 2 , depending on whether the latent invariant part C is fully informative about label Y , i.e., ( S, E ) ⊥⊥ Y | C .
Formal definitions of the corresponding SCMs are given as follows, where noises are omitted for simplicity [74, 77].
Assumption
[Equation]
Assumption
[Equation]
In the two SCMs above, f gen corresponds to the graph generation process in Assumption 2.1, and f spu is the mechanism describing how S is affected by C and E at the latent space.
By definition, S is directly controlled by C in FIIF and indirectly controlled by C through Y in PIIF, which can exhibit different behaviors in the observed distribution shifts.
In practice, performances of OOD algorithms can degrade dramatically if one of FIIF or PIIF is excluded [5, 71].
This issue can be more serious in graphs, since different distribution shifts can have different interaction modes at the latent space.
Moreover, f inv : C → Y indicates the labelling process, which assigns labels Y for the corresponding
G merely based on C .
Consequently, C is better clustered than S when given Y [13, 15, 86, 87], which also serves as the necessary separation assumption for a classification task [69, 16, 65].
Assumption 2.4 (Better Clustered Invariant Features) .
H ( C | Y ) ≤ H ( S | Y ) .
2.3 Challenges of OOD Generalization on Graphs
Built upon the graph generation process, we can formally derive the desired GNN that is able to generalize to OOD graphs under different distribution shifts, which implies the invariant GNN below 3 .
Definition 2.5 (Invariant GNN) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, considering a GNN ρ ◦ h that has a permutation invariant graph encoder h : G → R h and a downstream classifier ρ : R h → Y , ρ ◦ h is an invariant
GNN if it minimizes the worst case risk among all environments, i.e., min max e ∈E all R e .
Can existing methods produce a desired invariant GNN model?
We find the answers to be negative unfortunately.
Based on the synthetic BAMotif graph classification task [58, 104] shown in Fig.
Training Data
Testing Data “House” “House” “Cycle”
[Figure]
0.9 Data Biases
A cc u r a c y
Struc-ERMMixed-ERMStruc-IRMMixed-IRM
[Figure]
shifts.
0.9 Data Biases
A cc u r a c y
Struc-ERMMixed-ERMStruc-IRMMixed-IRM
[Figure]
we theoretically and empirically analyze whether existing methods could produce an invariant GNN, through the investigation of the following aspects.
More details and results are given in Appendix D.
Can GNNs trained with ERM generalize to OOD graphs?
As shown in Fig.
3, we find that GNNs trained with the standard empirical risk minimization (ERM) algorithm [97] are not able to generalize to OOD graphs.
As the data biases grows stronger, the performances of GNNs drop dramatically.
Furthermore, when graph size shifts are mixed in the data, GNNs can have larger variance at low data biases, indicating the instability of learning the desired relationships for the task.
The reason is that
ERM tends to overfit to the shortcuts or spurious correlations presented in specific substructures or attributes in the graphs [33].
This phenomenon has also been shown to exist in GNNs equipped with more sophisticated architectures such as attention mechanisms [99], under graph size shifts [46].
Can OOD objectives improve OOD generalization of GNNs?
Meanwhile, as shown in
Fig.
3, OOD objectives primarily developed on Euclidean data such as invariant risk minimization (IRM) [4] also cannot alleviate the problem.
On the contrary, IRM can fail catastrophically at non-linear regime if without sufficient support overlap for the test environments, i.e., ∪ e ∈E te supp ( P e ) ̸⊆ ∪ e ∈E tr supp ( P e ) [80].
In addition to IRM, the failure would also happen for alternative objectives [49, 9, 2] as proved by Rosenfeld et al.
Besides, different distribution shifts on graphs can be nested with each other where each one can have distinct spurious correlation type, e.g., FIIF or PIIF.
OOD objectives will also fail seriously if either of the correlation types is not supported [5, 71].
Moreover, non-trivial environment partitions or labels are required for performance guarantee of these OOD objectives [4, 49, 81, 2].
However, collecting meaningful environment partitions of graphs requires expert knowledge about graph data.
Thus, the environment labels can be expensive to obtain and are usually not available [67, 28, 37].
Alternative options such as random partitions tend not to alleviate the issue [23, 55], as it can be trivially deemed as mini-batching.
Challenges of OOD generalization on graphs.
The aforementioned failure analysis reveals that existing methods or objectives fail to elicit an invariant GNN primarily due to the following two challenges:
a) Distribution
shifts on graphs are more complicated where different types of spurious correlations can be entangled via different graph properties; b) Environment labels are usually not available due to the abstraction of graphs.
Despite these challenges, we are still highly motivated to address the following research question: Would it be possible to learn an invariant GNN that is generalizable under various distribution shifts by lifting the invariance principle to the graph data?
Invariance Principle for OOD Generalization on Graphs
We provide affirmative answers to the previous question by proposing a new framework, CIGA:
C ausality Inspired I nvariant G raph Le A rning.
Specifically, built upon the SCMs in Sec. 2.2, we generalize the invariance principle to graphs and instantiate the principle with theoretical guarantees.
3.1 Invariance for OOD Generalization on Graphs
Towards extending the invariance principle to graphs under SCMs in Sec. 2.2, we need to identify a set of variables that have stable causal relationship with Y under both FIIF and PIIF (Assumption 2.2, 2.3).
According to the ICM assumption [77], the labeling process C
→ Y is not informed nor influenced by other processes, implying that the conditional distribution P ( Y | C ) remains invariant to the interventions on the environment latent variable E [74].
Consequently, for a GNN with a permutation invariant encoder h : G → R h and a downstream classifier ρ : R h
→ Y , if h can recover the information of C from G in the learned graph representations, then the learning of ρ resembles traditional ERM [97] and can achieve the desired min-max optimality required by an invariant GNN (Def. 2.5).
However, recovering C from G is particularly difficult, since the generation of G from C involves two causal mechanisms f G c gen and f G gen in Assumption 2.1.
The unavailability of E further adds up the difficulty of enforcing the independence between the learned representations and E .
3.2 Invariant Graph Learning Framework
Causal algorithmic alignment.
To enable a GNN to learn to extract the information about C from
G , we propose the CIGA framework that explicitly aligns with the two causal mechanisms f G c gen and f G gen in Assumption 2.1.
The idea of alignment in CIGA is motivated by the algorithmic reasoning results that a neural network can learn a reasoning process better if its computation structure aligns with the process better [108, 110].
Specifically, we realize the alignment by decomposing a GNN into two sub-components 4 : a) a featurizer GNN g : G → G c aiming to identify the desired G c ; b) a classifier GNN f c : G c
→
Y that predicts the label Y based on the estimated G c , where G c refers to the space of subgraphs of G .
Formally, the learning objectives of f c and g can be formulated as:
[Equation]
where maximizing I ( (cid:98) G c ; Y ) is equivalent to minimizing a variational upper bound of R ( f c ( (cid:98) G c ))
[3, 120] that takes (cid:98) G c as inputs to predict label Y for G through f c and g , and (cid:98) G c is the estimated subgraph containing the information about C and hence needs to be independent of E .
Moreover, the extracted G c can either shares the same graph space with input G or has its own space with latent node and edge features, depending on the specific graph generation process.
In practice, architectures from the literature of interpretable GNNs are compatible with CIGA [122], hence can serve as practical choices for the implementation of CIGA.
More details are given in Appendix F.
Although we can technically align with the two causal mechanisms with g and f c , trivially optimizing this architecture cannot satisfy (cid:98) G c ⊥⊥ E .
Formally, merely maximizing I ( (cid:98) G c ; Y ) may include a subgraph from G s in (cid:98) G c since G s also shares certain mutual information with Y .
Moreover, the unavailability of E prevents the direct usage of E in enforcing the independence that is often adopted by previous methods [4, 49, 81, 31, 93], making the identification of G c more challenging.
Optimization objective.
To mitigate this issue, we need to find and translate other properties of G c into some differentiable and equivalent objectives to satisfy the independence constraint (cid:98) G c ⊥⊥ E .
The goal of the desired objective.
We begin by considering a simplistic setting where all the invariant subgraphs G c have the same size s c , i.e., | G c | =
s c 5 .
When maximizing I ( (cid:98) G c ; Y ) in Eq.
1, both FIIF and PIIF can introduce part of G s into (cid:98) G c .
In FIIF (Fig. 2(b)), as G c already contains the maximal possible information in G about Y , G c is a solution to max I ( (cid:98) G c ; Y ) .
However, some subgraph of
G c can be replaced by some subgraph of G s that is equally informative about Y .
In PIIF (Fig. 2(c)), there also exists some subgraph of G s that contains additional information about Y than G c , hence (cid:98) G c is more likely to involve some subgraph of G s .
Thus, the new objective needs to eliminate the auxiliary subgraphs of (cid:98) G c from G s such that the estimated (cid:98) G c can only contain G c .
An important property of G c .
Under both FIIF and PIIF SCMs (Fig. 2), for G e 1 c , G e 2 c that relate to the same causal factor c under two environments e 1 and e 2 , the desired (cid:98) G e 1 c , (cid:98) G e 2 c in e 1 and etend to have high mutual information, i.e., ( G e 1 c , G e 2 c ) ∈ arg max I ( (cid:98) G e 1 c ; (cid:98) G e 2 c ) .
While for G e 1 c and another G e 1 c ′ corresponding to a different c ′ ̸
= c , under the same environment e 1 , including any subgraph from G e 1 s in (cid:98) G e 1 c , (cid:98) G e 1 c ′ will enlarge their mutual information, or in other words, ( G e 1 c , G e 1 c ′ ) ∈ arg min I ( (cid:98) G e 1 c ; (cid:98) G e 1 c ′ ) .
Thus, we can derive an important property of G c , that is, ∀ e 1 , e 2 ∈ E all ,
[Equation]
where (cid:98) G e 1 c
and (cid:98) G e 2 c are the estimated invariant subgraphs corresponding to the same causal factor c under environment e
and e 2 , respectively, while (cid:98) G e 2 c ′ corresponds to a different causal factor c ′ .
Deriving CIGA v1 based on the identified property of G c .
In practice, C is not given.
Nevertheless, since C and Y shares a stable causal relationship in both FIIF and PIIF SCMs, Y can serve as a proxy of C in Eq. 2.
Moreover, as Eq. 2 holds for any ∀ e 1 , e 2 ∈ E all , the environment superscripts can be eliminated without affecting Eq. 2.
Furthermore, when both I ( (cid:98) G e 1 c ; (cid:98) G e 2 c | C = c ) and I ( (cid:98) G c ; Y ) are maximized, I ( (cid:98) G e 1 c ; (cid:98) G e 1 c ′ | C
=
c
′ , c ′ ̸
= c ) is automatically minimized, otherwise all classes will collapse to trivial solutions which is contradictory given I ( (cid:98) G c ; Y ) being maximized.
Therefore, we can derive an alternative objective to Eq.
1 by leveraging Eq.
2 to replace the independence condition: ( CIGAv1 )
[Equation]
where (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G is sampled from training graphs that share the same label Y as G .
In Theorem 3.1, we show how Eq. 3 is equivalent to Eq.
Nevertheless, Eq. 3 requires a strong assumption on the size of G c .
However, the size of G c is usually unknown or changes for different C s. In this circumstance, maximizing Eq.
2 without additional constraints will lead to the presence of part of G s in (cid:98) G c .
For instance, (cid:98) G c =
G is a trivial solution to Eq.
3 when s c = ∞ .
Deriving CIGA v2 by resolving size constraint on G c in CIGA v1.
To this end, we further resort to the properties of G s .
In both FIIF and PIIF SCMs (Fig. 2), G s and G c can share certain overlapped information about Y .
When maximizing I ( (cid:98) G c ; (cid:101) G c | Y ) and I ( (cid:98) G c ; Y ) , the appearance of partial G s in (cid:98) G c will not affect the optimality.
However, it can reduce the mutual information between the left part (cid:98) G s =
G − (cid:98) G c and Y , i.e., I ( (cid:98) G s ; Y ) .
Therefore, by maximizing I ( (cid:98) G s ; Y ) , we can reduce including part of G s into (cid:98) G c .
Meanwhile, to avoid trivial solution that G c ⊆ (cid:98) G s during maximizing
I ( (cid:98) G s ; Y ) , we can leverage the better clustering property of G c implied by Assumption 2.4 to derive the constraint I ( (cid:98) G s ; Y ) ≤
I ( (cid:98) G c ; Y ) .
Thus, we can obtain a new objective CIGAv2 as follows:
[Equation]
where (cid:98) G c
=
g ( G ) , (cid:101) G c
=
g ( (cid:101) G ) and (cid:101) G ∼ P ( G | Y ) , i.e., (cid:101) G is sampled from training graphs that share the same label Y as G .
We also prove the equivalence between Eq.
4 and Eq. 1 in Theorem 3.1.
3.3 Theoretical Analysis and Practical Discussions
Theorem 3.1 (CIGA Induces Invariant GNNs) .
Given a set of graph datasets {D e } e and environments E all that follow the same graph generation process in Sec. 2.2, assuming that (a) f G gen and f G c gen in Assumption 2.1 are invertible, (b) samples from each training environment are equally distributed, i.e., |D ˆ e | = |D
˜ e | , ∀ ˆ e, ˜ e ∈ E tr , then:
(i).
If ∀ G c , | G c | = s c , then each solution to Eq. 3, elicits an invariant GNN (Def. 2.5).
(ii).
Each solution to Eq. 4, elicits an invariant GNN (Def. 2.5).
We prove Theorem 3.1 (i) and (ii) in Appendix E.2, E.3, respectively.
Practical implementations of CIGA objectives.
After showing the power of CIGA, we introduce the practical implementations of CIGAv1 and CIGAv2 objectives.
Specifically, an exact estimate of the second term I ( (cid:98) G c ; (cid:101) G c | Y ) could be highly expensive [96, 8].
However, contrastive learning with supervised sampling provides a practical solution for the approximation [42, 20, 82, 96, 8]:
[Equation] [Equation]
[Table]
where positive samples ( (cid:98) G c , (cid:101) G c ) are the extracted subgraphs of graphs that share the same label as
G , negative samples are those having different labels, P g ( G |Y =
Y ) is the push-forward distribution of P ( G |Y =
Y ) by featurizer g , P ( G |Y =
Y ) refers to the distribution of G given the label Y , P ( G
|Y ̸ = Y ) refers to the distribution of G given the label that is different from Y , h (cid:98) G c , h (cid:101) G c , h G ic are the graph presentations of the estimated subgraphs, and ϕ is the similarity metric for graph representations.
As M → ∞ , Eq.
5 approximates I ( (cid:98) G c ; (cid:101) G c | Y ) , which can be regarded as a nonparameteric resubstitution entropy estimator via the von Mises-Fisher kernel density [1, 41, 101].
Thus, plugging it into Eq. 3 and Eq.
4 can relieve the issue of approximating I ( (cid:98) G c ; (cid:101) G c | Y ) in practice.
To implement I ( (cid:98) G s ; Y ) given the constraint I ( (cid:98) G s ; Y )
≤
I ( (cid:98) G c ; Y ) in CIGAv2, a practical choice is to adopt hinge loss that implement the constrained I ( (cid:98) G s ; Y ) as 1 N R (cid:98) G s · I ( R (cid:98) G c ≤ R (cid:98) G s ) , where N is the number of samples, I is an indicator function that outputs 1 when the inner condition is satisfied otherwise 0 , and R (cid:98) G s and R (cid:98) G c are the empirical risk vector of the predictions for each sample based on the corresponding (cid:98) G s and (cid:98) G c .
More implementation details can be found in Appendix F.
Discussions and implications of CIGA.
Although using contrastive learning to improve OOD generalization is not new in the literature [27, 61, 124], previous methods cannot yield OOD guarantees in graph circumstances due to the highly non-linearity and the unavailability of domain labels
E .
In particular, CIGA can be reduced to directly applying contrastive learning when without the decomposition for causal algorithmic alignment.
However, in the experiments we found that merely using the contrastive objective, i.e., CNC [124], yields unsatisfactory OOD generalization performance, which further implies the necessity of the decomposition in CIGA.
Moreover, the architecture of CIGA can have multiple other implementations for both the featurizer and classifier, such as identifying G c at the latent space [86, 87].
Since we cannot enumerate every possible implementation, in this work we choose interpretable GNN architectures as a prototype validation for CIGA and leave more sophisticated architectures as future works.
In particular, when optimized with ERM objective, CIGA can be reduced to interpretable GNNs .
However, merely using interpretable GNNs such as ASAP [78], GIB [120] or DIR [104] cannot yield satisfactory OOD performance.
As shown in Table 1(b) and discussed in Appendix.
D.4, GIB can only work for FIIF, while DIR cannot yield OOD guarantees for neither FIIF and PIIF SCMs.
These results are also empirically validated in the experiments.
We provide more detailed discussions in Appendix B.
Empirical Studies
We conduct extensive experiments with 16 datasets to verify the effectiveness of CIGA.
Datasets.
We use the SPMotif datasets from DIR [104] where artificial structural shifts and graph size shifts are nested (SPMotif-Struc).
Besides, we construct a harder version mixed with attribute shifts (SPMotif-Mixed).
To examine CIGA in real-world scenarios with more complicated relationships and distribution shifts, we also use DrugOOD [40] from AI-aided Drug Discovery with Assay,
Scaffold, and Size splits, convert the ColoredMNIST from IRM [4] using the algorithm from Knyazev et al.
[46] to inject attribute shifts, and split Graph-SST [122] to inject degree biases.
To compare with previous specialized OOD methods for graph size shifts [113, 11], we use the datasets in Bevilacqua et al.
[11] that are converted from TU benchmarks [67].
More details can be found in Appendix G.1.
Baselines and our methods.
Besides the ERM, we also compare with SOTA interpretable GNNs,
GIB [120], ASAP Pooling [78], and DIR [104], to validate the effectiveness of the optimization
[Table]
objective in CIGA.
We use the same selection ratio (i
[Table]
e., s c ) for all models.
Moreover, to validate the effectiveness of the decomposition in CIGA, we compare CIGA with SOTA OOD objectives including IRM [4], v-Rex [49] and IB-IRM [2], for which we apply random environment partitions following [23].
We also compare CIGA with EIIL [23] and CNC [124] that do not require environment labels, where CNC [124] has a more sophisticated contrastive sampling strategy for combating subpopulation shifts.
More implementation and comparison details are deferred to Appendix G.2.
Evaluation.
We report the classification accuracy for all datasets, except for DrugOOD datasets where we use ROC-AUC following [40], and for TU datasets where we use Matthews correlation coefficient following [11].
We repeat the evaluation multiple times, select models based on the validation performances, and report the mean and standard deviation of the corresponding metric.
For each dataset, we also report the “Oracle” performances that run ERM on the randomly shuffled data.
OOD generalization performance on structure and mixed shifts.
In Table 1, we report the test accuracy of each method, where we omit GIB due to its poor convergence.
Different biases indicate different strengths of the distribution shifts.
Although the training accuracy of most methods converges to more than 99% , the test accuracy decreases dramatically as the bias increases and as more distribution shifts are mixed, which concurs with our discussions in Sec.
2.3 and Appendix D.
Due to the simplicity of the task as well as the relatively high support overlap between training and test distributions, interpretable GNNs and OOD objectives can improve certain OOD performance, while they can have high variance since they donot have OOD generalization guarantees.
In contrast,
CIGAv1 and CIGAv2 outperform all of the baselines by a significant margin up to 10% with lower variance , which demonstrates the effectiveness and excellent OOD generalization ability of CIGA.
[Table]
OOD generalization performance on realistic shifts.
In
Table 2 and Table 3, we examine the effectiveness of CIGA in real-world data and more complicated distribution shifts. Both averaged accuracy and ranks are reported because of the dataset heterogeneity.
Since the tasks are harder than synthetic ones, interpretable GNNs and OOD objectives perform similar to or even under-perform the ERM baselines, which is also consistent to the observations in non-linear benchmarks [34, 40].
However, both CIGAv1 and CIGAv2 consistently and significantly outperform previous methods, including previous specialized methods Γ
GNNs [11] for combating graph size shifts, demonstrating the generality and superiority of CIGA. Comparisons with advanced ablation variants.
As discussed in Sec. 3.3, CIGA can be reduced to interpretable GNNs and contrastive learning approaches. However, across all experiments, we can observe that neither the advanced interpretable GNNs (DIR) nor sophisticated contrastive objectives with specialized sampling strategy (CNC) can yield satisfactory OOD performance, which serves as strong evidence for the necessities of the decomposition as well as the objective in CIGA. Furthermore, although CIGAv1 can outperform CIGAv2 when we may have a relatively accurate
[Figure]
t e s t a u c ermCIGAv
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
t e s t a u c ermCIGAv
[Figure]
m
a tt h e w s c o rr e l a t i o n c o e ff i
c i e n t ermCIGAv
[Figure]
[Figure]
i c
[Figure]
s c , the improvements in CIGAv1 are not as stable as CIGAv2 or even unsatisfactory when the assumption is violated.
This phenomenon also reveals the superiority of CIGAv2 in practice.
Hyperparameter sensitivity analysis.
To examine how sensitive CIGA is to the hyperparamters α and β for contrastive loss and hinge loss, respectively.
We conduct experiments based on the hardest datasets from each table (i.e., SPMotif-Mixed with the bias of 0 .
9 , DrugOOD-Scaffold and the NCI109 datasets from Table 1, Table 2, and Table 3, respectively.)
with different α and β .
When changing the value of β , we fix the α to a specific value under which the model has a relatively good performance (but not the best, to fully examine the robustness of CIGA in practice).
The results are shown in Fig.
4 and Fig.
It can be found that both CIGAv1 and CIGAv2 are robust to different values of α and β , respectively, across different datasets and distribution shifts.
Besides, the results also reflect the effects of the additional penalty terms in CIGA.
For example, in Fig.
16, when α is too small, the invariance of the identified invariant subgraphs (cid:98) G c may not be guaranteed, resulting worse performances.
Similarly, as shown in Fig. 17, when β becomes too small, some part of the spurious subgraph may still appear in the estimated invariant subgraphs, which yields worse performances.
Besides, when α and β become too large, the optimization of CIGA can be affected due to their intrinsic conflicts with ERM, hence a better optimization scheme for CIGA can be a promising future direction [18].
We provide more details and additional analysis on the efficiency of
CIGA and single environment
OOD generalization performance of CIGA in Appendix G.4, as well as the visualization examples of the identified invariant subgraph in Appendix G.5.
Conclusions
We studied the OOD generalization on graphs via graph classification, and propose a new solution
CIGA through the lens of causality.
By modeling potential distribution shifts on graphs with SCMs, we generalized and instantiated the invariance principle to graphs, which was shown to have promising theoretical and empirical OOD generalization ability under a variety of distribution shifts.
Acknowledgments and Disclosure of Funding
We thank the reviewers for their valuable comments.
This work was supported by GRFfrom the RGC of HKSAR and CUHK direct grant 4055146.
TL was partially supported by Australian
Research Council Projects DP180103424, DE-190101473, IC-190100031, DP-220102121, and
FT-220100318.
YZ and BH were supported by the RGC Early Career Scheme No. 22200720, NSFC
Young Scientists Fund No. 62006202, Guangdong Basic and Applied Basic Research Foundation No.
2022A1515011652, and Tencent AI Lab Rhino-Bird Gift Fund.