SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yongqiang Chen $^{1*}$ , Yatao Bian $^{2}$ , Kaiwen Zhou $^{1}$
<sup>1</sup>The Chinese University of Hong Kong <sup>2</sup>Tencent AI Lab
{yqchen, kwzhou}@cse.cuhk.edu.hk yatao.bian@gmail.com
Binghui Xie $^{1}$ , Bo Han $^{3}$ , James Cheng $^{1}$
$^{3}$ Hong Kong Baptist University
bhanml@comp.hkbu.edu.hk {bhxie21,jcheng}@cse.cuhk.edu.hk
Invariant graph representation learning aims to learn the invariance among data from different environments for out-of-distribution generalization on graphs. As the graph environment partitions are usually expensive to obtain, augmenting the environment information has become the de facto approach. However, the usefulness of the augmented environment information has never been verified. In this work, we find that it is fundamentally impossible to learn invariant graph representations via environment augmentation without additional assumptions. Therefore, we develop a set of minimal assumptions, including variation sufficiency and variation consistency, for feasible invariant graph learning. We then propose a new framework Graph invArient Learning Assistant (GALA). GALA incorporates an assistant model that needs to be sensitive to graph environment changes or distribution shifts. The correctness of the proxy predictions by the assistant model hence can differentiate the variations in spurious subgraphs. We show that extracting the maximally invariant subgraph to the proxy predictions provably identifies the underlying invariant subgraph for successful OOD generalization under the established minimal assumptions. Extensive experiments on 12 datasets including DrugOOD with various graph distribution shifts confirm the effectiveness of GALA<sup>1</sup>.
Graph representation learning with graph neural networks (GNNs) has proven to be highly successful in tasks involving relational information [24, 35, 68, 74, 75]. However, it assumes that the training and test graphs are independently drawn from the identical distribution (iid.), which can hardly hold for many graph applications such as in Social Network, and Drug Discovery [25, 26, 37, 70, 86, 94]. The performance of GNNs could be seriously degenerated by graph distribution shifts, i.e., mismatches between the training and test graph distributions caused by some underlying environmental factors during the graph data collection process [16, 21, 28, 70, 87, 94]. To overcome the Out-of-Distribution (OOD) generalization failure, recently there has been a growing surge of interest in incorporating the invariance principle from causality [60] into GNNs [8, 17, 41, 45, 55, 72, 73, 77, 83]. The rationale of the invariant graph learning approaches is to identify the underlying invariant subgraph of the input graph, which shares an invariant correlation with the target labels across multiple graph
*Work done during an internship at Tencent AI Lab.
$^{1}$ Code is available at https://github.com/LFphase/GALA.
Figure 1: An illustration of GALA with the task of classifying graphs according to whether there exists a "House" or "Cycle" motif. Given the training data where the "House" subgraph often co-occurs with a "Grid" and the "Cycle" subgraph often co-occurs with a "Hexagon". An ERM trained environment assistant model will fit the spurious subgraph and therefore yield proxy predictions "House" or "Cycle" for any graphs containing a "Grid" (left half) or "Hexagon" (right half), respectively. GALA first separates the samples according to the correctness of the proxy predictions into the sets of positive graphs  $\{G^p\}$  (correct, in blue) and negative graphs  $\{G^n\}$  (incorrect, in green). Then, GALA extracts the maximally invariant subgraph among  $\{G^p\}$  and  $\{G^n\}$ , i.e., pulling graphs with the same graph label but from  $\{G^p\}$  and  $\{G^n\}$  closer in the latent space, hence identifies the invariant subgraph.
distributions from different environments [8, 72]. Thus, the predictions made merely based on the invariant subgraphs can be generalized to OOD graphs that come from a new environment [60].
As the environment labels or partitions on graphs are often expensive to obtain [8], augmenting the environment information, such as generating new environments [45, 72, 73] and inferring the environment labels [41, 77], has become the de facto approach for invariant graph learning. However, little attention has been paid to verifying the fidelity (or faithfulness²) of the augmented environment information. For example, if the generated environments or inferred environment labels induce a higher bias or noise, it would make the learning of graph invariance even harder. Although it looks appealing to learn both the environment information and the graph invariance, the existing approaches could easily run into the "no free lunch" dilemma [71]. In fact, Lin et al. [43] found that there exist negative cases in the Euclidean regime where it is impossible to identify the invariant features without environment partitions. When it comes to the graph regime where the OOD generalization is fundamentally more difficult [8] than the Euclidean regime, it raises a challenging research question:
When and how could one learn graph invariance without the environment labels?
In this work, we present a theoretical investigation of the problem and seek a set of minimal assumptions on the underlying environments for feasible invariant graph learning. Based on a family of simple graph examples (Def. 3.1), we show that existing environment generation approaches can fail to generate faithful environments, when the underlying environments are not sufficient to uncover all the variations of the spurious subgraphs (Prop. 3.2). On the contrary, incorporating the generated environments can even lead to a worse OOD performance. The failure of faithful environment generation implies the necessity of variation sufficiency (Assumption 3.3). Moreover, even with sufficient environments, inferring faithful environment labels remains impossible. Since invariant and spurious subgraphs can have an arbitrary degree of correlation with labels, there exist multiple sets of training environments that have the same joint distribution of  $P(G,Y)$  but different invariant subgraphs. Any invariant graph learning algorithms will inevitably fail to identify the invariant subgraph in at least one set of training environments (Prop. 3.4). Therefore, we need to additionally ensure the variation consistency (Assumption 3.5), that is, the invariant and spurious subgraphs should have a consistent relationship in the correlation strengths with the labels.
To resolve the OOD generalization challenge under the established assumptions, we propose a new framework Graph invArient Learning Assistant (GALA). GALA incorporates an additional assistant model that needs to be prone to distribution shifts, to generate proxy predictions of the training samples. Different from previous environment inferring approaches [41, 77], GALA does not require explicit environment labels but merely proxy predictions to differentiate the variations in the spurious
2The fidelity or faithfulness refers to whether the augmented environment information can actually improve the OOD generalization on graphs.
subgraphs. As shown in Fig. 1, we first fit an environment assistant model to the training distribution and then divide the training graphs into a positive set  $\{G^p\}$  and a negative  $\{G^n\}$ , according to whether the proxy predictions are correct or not, respectively. As spurious correlations tend to vary more easily than invariant correlations, the variations in spurious subgraphs are further differentiated and increased between  $\{G^p\}$  and  $\{G^n\}$ . Then, only the invariant subgraph holds an invariant correlation with the label among  $\{G^p\}$  and  $\{G^n\}$ , and hence can be identified by extracting the subgraphs that maximize the intra-class subgraph mutual information among  $\{G^p\}$  and  $\{G^n\}$  (Theorem 4.1).
We conduct extensive experiments to validate the effectiveness of GALA using 12 datasets with various graph distribution shifts. Notably, GALA brings improvements up to  $30\%$  in multiple graph datasets.
- We identify failure cases of existing invariant graph learning approaches and establish the minimal assumptions for feasible invariant graph learning;
- We develop a novel framework GALA with provable identifiability of the invariant subgraph for OOD generalization on graphs under the assumptions;
- We conduct extensive experiments to verify both our theoretical results and the superiority of GALA;
Notably, both our theory and solution differ from Lin et al. [43] fundamentally, as we do not rely on the auxiliary information and are compatible with the existing interpretable and generalizable GNN architecture for OOD generalization on graphs. Meanwhile, we provide a new theoretical framework that resolves the counterexample in Lin et al. [43] while enjoying provable identifiability.
2 Background and Preliminaries
We begin by introducing the key concepts and backgrounds of invariant graph learning, and leave more details in Appendix C. The notations used in the paper are given in Appendix A.
OD generalization on graphs. This work focuses on graph classification, while the results generalize to node classification as well using the same setting as in Wu et al. [72]. Specifically, we are given a set of graph datasets  $\mathcal{D} = \{\mathcal{D}_e\}_{e\in \mathcal{E}_{\mathrm{all}}}$  collected from multiple environments  $\mathcal{E}_{\mathrm{all}}$ . Samples  $(G_i^e,Y_i^e)\in \mathcal{D}^e$  from the environment  $e$  are drawn independently from an identical distribution  $\mathbb{P}^e$ . The goal of OOD generalization on graphs is to find a GNN  $f$  that
(a)  $\mathcal{G}$ -Gen. SCM
(b) FIIF SCM
(c) PIIF SCM
Figure 2: SCMs on graph distribution shifts [8].
minimizes the maximal loss among all environments, i.e., to minimize  $\max_{e\in \mathcal{E}_{\mathrm{all}}}R^{e}$ , where  $R^{e}$  is the risk of  $f$  under environment  $e$ . We consider the same graph generation process proposed by Chen et al. [8] which is inspired by real-world drug discovery task [57] and covers a broad case of graph distribution shifts. As shown in Fig. 2, the generation of the observed graphs  $G$  and labels  $Y$  are controlled by a latent causal variable  $C$  and a spurious variable  $S$ .  $C$  and  $S$  control  $Y$  and  $G$  by controlling the generation of the underlying invariant subgraph  $G_{c}$  and spurious subgraph  $G_{s}$ , respectively. Since  $S$  can be affected by the environment  $E$ , the correlation between  $Y$  and  $G_{s}$  can change arbitrarily when the environment changes. Besides, the interaction among  $C$ ,  $S$  and  $Y$  at the latent space can be further categorized into Full Informative Invariant Features (FIIF) when  $Y\perp S|C$ , and Partially Informative Invariant Features (PIIF) when  $Y\not\perp S|C$ .
To tackle the OOD generalization challenge on graphs from Fig. 2, the existing invariant graph learning approaches are generically designed to identify the underlying invariant subgraph  $G_{c}$  to predict the label  $Y$  [8, 72]. Specifically, the goal of OOD generalization on graphs is to learn an invariant GNN  $f \coloneqq f_{c} \circ g$ , which is composed of: a) a featurizer  $g: \mathcal{G} \to \mathcal{G}_{c}$  that estimates the invariant subgraph  $\widehat{G}_{c}$ ; b) a classifier  $f_{c}: \mathcal{G}_{c} \to \mathcal{V}$  that predicts the label  $Y$  based on the extracted  $\widehat{G}_{c}$ , where  $\mathcal{G}_{c}$  refers to the space of subgraphs of  $\mathcal{G}$ . The learning objectives of  $f_{c}$  and  $g$  are formulated as
$$
\max  _ {f _ {c}, g} I \left(\widehat {G} _ {c}; Y\right), \text {s . t .} \widehat {G} _ {c} \perp E, \widehat {G} _ {c} = g (G). \tag {1}
$$
Since  $E$  is not observed, many strategies are proposed to impose the independence of  $\widehat{G}_c$  and  $E$ . A prevalent approach is to augment the environment information. Based on the estimated invariant
subgraphs  $\widehat{G}_c$  and spurious subgraphs  $\widehat{G}_s$ , Liu et al. [45], Wu et al. [72, 73] propose to generate new environments, while Li et al. [41], Yang et al. [77] propose to infer the underlying environment labels. However, we show that they all fail to augment faithful environment information in Sec. 3.
Besides, Miao et al. [55, 56], Yu et al. [81, 82, 83] adopt graph information bottleneck to tackle FIIF graph shifts, but they cannot generalize to PIIF shifts, while Our work focuses on PIIF shifts as it is more challenging when without environment labels [43]. Fan et al. [17] generalize [40] to tackle severe graph biases, i.e., when  $H(S|Y) < H(C|Y)$ . Chen et al. [8] propose a contrastive framework to tackle both FIIF and PIFF graph shifts, but is limited to  $H(S|Y) > H(C|Y)$ . In practice, as it is usually unknown which correlation is stronger, we need a unified solution to tackle both cases.
Invariant learning without environment labels. In the Euclidean regime, there are plentiful studies in invariant learning without environment labels. Creager et al. [12] propose a minmax formulation to infer the environment labels. Liu et al. [46] propose a self-boosting framework based on the estimated invariant and variant features. Deng et al. [14], Liu et al. [44], Pezeshki et al. [61], Zhang et al. [85] propose to infer labels based on the failures of an ERM model. However, Lin et al. [43] find failure cases of the aforementioned approaches that it is impossible to identify the invariant features without given environment labels in Euclidean data, and propose a solution that leverages auxiliary environment information for invariant learning. As the OOD generalization on graphs poses more challenges [8], whether it is feasible to learn invariant graph representations without any auxiliary environment information remains elusive.
3 Pitfalls of Environment Augmentation
Given only the mixed training data without environment partitions, is it possible to learn to generate faithful environments or infer the underlying environment labels that facilitate OOD generalization on graphs? In the discussion below, we adopt the two-piece graphs to instantiate the problem, which is the simplistic version of the PIIF distribution shifts in Fig. 2(c), motivated by Kamath et al. [30].
Definition 3.1 (Two-piece graphs). Each environment  $e$  is defined with two parameters,  $\alpha_{e}, \beta_{e} \in [0,1]$ , and the dataset  $(G^{e}, Y^{e}) \in \mathcal{D}_{e}$  is generated as follows:
(a) Sample  $Y^{e} \in \{-1, 1\}$  uniformly;
(b) Generate  $G_{c}$  and  $G_{s}$  via:  $G_{c} := f_{gen}^{G_{c}}(Y^{e} \cdot \text{Rad}(\alpha_{e}))$ ,  $G_{s} := f_{gen}^{G_{s}}(Y^{e} \cdot \text{Rad}(\beta_{e}))$ , where  $f_{gen}^{G_{c}}, f_{gen}^{G_{s}}$  map the input  $\{-1,1\}$  to a corresponding graph selected from a given set, and  $\text{Rad}(\alpha)$  is a random variable taking value  $-1$  with probability  $\alpha$  and  $+1$  with  $1 - \alpha$ ;
(c) Synthesize  $G^{e}$  by randomly assembling  $G_{c}$  and  $G_{s}$ :  $G^{e} \coloneqq f_{gen}^{G}(G_{c}, G_{s})$ .
We denote an environment  $e$  with  $(\alpha, \beta_e)$  for simplicity. Different environments will have a different  $\beta_e$ , thus  $P(Y|G_s)$  will change across different environments, while  $P(Y|G_c)$  remains invariant.
3.1 Pitfalls of environment generation
We begin by discussing the cases where there are few environments, and generating new environments is necessary [45, 72, 73]. Environment generation aims to provide some additional "virtual" environments  $\mathcal{E}_v$  such that the invariant subgraph can be identified via applying an OOD risk to the joint dataset with the augmented data  $\mathcal{D}_{\mathrm{tr}}^v = \{\mathcal{D}_e|e\in \mathcal{E}_{\mathrm{tr}}\cup \mathcal{E}_v\}$ .
The generation of "virtual" environments is primarily based on the intermediate estimation of the invariant and spurious subgraphs, denoted as  $\widehat{G}_c$  and  $\widehat{G}_s$ , respectively. Liu et al. [45], Wu et al. [73] propose DIR and GREA to construct new graphs by assembling  $\widehat{G}_c$  and  $\widehat{G}_s$  from different graphs. Specifically, given  $n$  samples  $\{G^i,Y^i\}_{i = 1}^n$ , the new graph samples in  $\mathcal{E}_v$  is generated as follows:
$$
G ^ {i, j} = f _ {\text {g e n}} ^ {G} \left(\widehat {G} _ {c} ^ {i}, \widehat {G} _ {s} ^ {j}\right), \forall i, j \in \{1 \dots n \}, Y ^ {i, j} = Y ^ {i},
$$
which generates a new environment  $\mathcal{E}_v$  with  $n^2$  samples. Although both DIR and GREA gain some empirical success, the faithfulness of  $\mathcal{E}_v$  remains questionable, as the generation is merely based on
<sup>3</sup>We slightly abuse the superscript and subscript when denoting the  $i$ th sample to avoid confusion of double superscripts or subscripts.
inaccurate estimations of the invariant and spurious subgraphs. Specifically, when  $\widehat{G}_c$  contains parts of  $G_s$ , assigning the same labels to the generated graph is more likely to strengthen the spurious correlation between  $G_s$  and  $Y$ . For example, when the model yields a reversed estimation, i.e.,  $\widehat{G}_c = G_s$  and  $\widehat{G}_s = G_c$ , the generated environment will destroy the invariant correlations.
Proposition 3.2. Consider the two-piece graph dataset  $\mathcal{E}_{tr} = \{(\alpha, \beta_1), (\alpha, \beta_2)\}$  with  $\alpha \geq \beta_1, \beta_2$  (e.g.,  $\mathcal{E}_{tr} = \{(0.25, 0.1), (0.25, 0.2)\}$ ), and its corresponding mixed environment  $\mathcal{E}_{tr}^{mix} = \{(\alpha, (\beta_1 + \beta_2)/2)\}$  (e.g.,  $\mathcal{E}_{tr}^{mix} = \{(0.25, 0.15)\}$ ). When  $\widehat{G}_c = G_s$  and  $\widehat{G}_s = G_c$ , it holds that the augmented environment  $\mathcal{E}_v$  is also a two-piece graph dataset with
$$
\mathcal {E} _ {v} = \left\{\left(0. 5, \left(\beta_ {1} + \beta_ {2}\right) / 2\right) \right\} (e. g., \mathcal {E} _ {v} = \{(0. 5, 0. 1 5) \}).
$$
The proof is given in Appendix E.1. This also extends to the adversarial augmentation [72, 83], which will destroy the actual  $\widehat{G}_c$ . As both DIR and GREA adopt the same environment generation procedure, we verify the failures of environment generation with GREA in Table 2 of Sec. 5, where GREA can perform comparably with ERM. In fact, when the underlying environments are insufficient to differentiate the variations of the spurious features, it is fundamentally impossible to identify the underlying invariant graph from the spurious subgraph. More formally, if  $\exists G_s$ , such that  $P^{e_1}(Y|G_s) = P^{e_2}(Y|G_s)$  for any  $e_1, e_2 \in \mathcal{E}_{\mathrm{tr}}$ , where  $P^e (Y|G_s)$  is the conditional distribution  $P(Y|G_{s})$  under environment  $e \in \mathcal{E}_{\mathrm{all}}$ , it is impossible for any graph learning algorithm to identify  $G_{c}$ . We provide a formal discussion in Appendix E.2. The failure implies a fundamental requirement that  $\mathcal{E}_{\mathrm{tr}}$  should uncover all the potential variations in the spurious subgraph.
Assumption 3.3. (Variation sufficiency) For graphs generated following Fig. 2, for any  $G_{s}$ ,  $\exists e_{1}, e_{2} \in \mathcal{E}_{tr}$ , such that  $P^{e_1}(Y|G_s) \neq P^{e_2}(Y|G_s)$ , and  $P^{e_1}(Y|G_c) = P^{e_2}(Y|G_c)$ .
Assumption 3.3 aligns with the definition of invariance [8, 30] that the invariant subgraph  $G_{c}$  is expected to satisfy  $P^{e_1}(Y|G_c) = P^{e_2}(Y|G_c)$  for  $e_1, e_2 \in \mathcal{E}_{\mathrm{all}}$ . If there exists  $G_{s}$  satisfying the invariance condition as well, then it is impossible to tell  $G_{c}$  from  $G_{s}$  even with environment labels.
3.2 Pitfalls of environment inferring
Although environment sufficiency (Assumption 3.3) relieves the need for generating new environments, is it possible to infer the underlying environment labels via approaches such as Mo1e00D [77] and GIL [41], to facilitate invariant graph learning? Unfortunately, we find a negative answer.
Considering the two-piece graph examples  $\mathcal{E}_{\mathrm{tr}} = \{(0.2, 0.1), (0.2, 0.3)\}$ , when given the underlying environment labels, it is easy to identify the invariant subgraphs from spurious subgraphs. However, when the environment labels are not available, we have the mixed data as  $\mathcal{E}_{\mathrm{tr}} = \{(0.2, 0.2)\}$ , where  $P(Y|G_c) = P(Y|G_s)$ . The identifiability of  $G_s$  is ill-posed, as it does not affect the  $\mathcal{E}_{\mathrm{tr}}$  even if we swap  $G_c$  and  $G_s$ . More formally, considering the environment mixed from two two-piece graph environments  $\{(\alpha, \beta_1)\}$  and  $\{(\alpha, \beta_2)\}$ , then we have  $\mathcal{E}_{\mathrm{tr}} = \{(\alpha, (\beta_1 + \beta_2)/2)\}$ . For each  $\mathcal{E}_{\mathrm{tr}}$ , we can also find a corresponding  $\mathcal{E}_{\mathrm{tr}}' = \{((\beta_1' + \beta_1')/2, \alpha')\}$  with  $\{(\beta_1', \alpha')\}$  and  $\{(\beta_2', \alpha')\}$ . Then, let
$$
\alpha = \left(\beta_ {1} ^ {\prime} + \beta_ {1} ^ {\prime}\right) / 2 = \alpha^ {\prime} = \left(\beta_ {1} + \beta_ {2}\right) / 2. \tag {2}
$$
We now obtain  $\mathcal{E}_{\mathrm{tr}}$  and  $\mathcal{E}_{\mathrm{tr}}'$  which share the same joint distribution  $P(Y, G)$  while the underlying  $G_{c}$  is completely different. More generally, we have the following proposition.
Proposition 3.4. There exist 2 two-piece graph training environments  $\mathcal{E}_{tr}$  and  $\mathcal{E}_{tr}^{\prime}$  that share the same joint distribution  $P(Y,G)$ . Any learning algorithm will fail in either  $\mathcal{E}_{tr}$  or  $\mathcal{E}_{tr}^{\prime}$ .
The proof is given in Appendix E.3. The experiments in Sec. 5 validate that both Mole00D and GIL fail to infer faithful environment labels and even underperform ERM. It implies that whenever it allows the existence of an identical training distribution by mixing the environments, invariant graph learning is impossible. Therefore, we need an additional assumption that excludes the unidentifiable case. We propose to constrain the relationship between  $\alpha$  (i.e.,  $H(Y|G_c)$ ) and  $\beta_{e}$  (i.e.,  $H(Y|G_s)$ ).
Assumption 3.5. (Variation consistency) For all environments in  $\mathcal{E}_{tr}$ ,  $H(C|Y) \neq H(S|Y)$ .
Intuitively, Assumption 3.5 imposes the consistency requirement on the correlation strengths between invariant and spurious subgraphs with labels. For two-piece graphs with consistent variations, mixing up the environments will yield a new environment with the same variation strength relationships.
Thus, Assumption 3.5 gets rid of the previous unidentifiable cases. Moreover, Assumption 3.5 also aligns with many realistic cases. For example, the relation of a specific functional group (e.g., -OH) with a molecule can hardly be reversed to that held upon the scaffold of the molecule, due to the data collection process. Therefore, Assumption 3.5 also resolves the counterexample proposed by Lin et al. [43]. Different from our work, Lin et al. [43] propose to incorporate additional auxiliary information that satisfies certain requirements to mitigate the unidentifiable case. However, such auxiliary information is often unavailable and expensive to obtain on graphs. More importantly, the requirements are also unverifiable without more assumptions, which motivates us to consider the relaxed case implied by Assumption 3.5.
3.3 Challenges of environment augmentation
To summarize, the two assumptions constitute the minimal assumptions for feasible invariant graph learning. Failing to satisfy either one of them while lacking additional inductive biases will result in the "no free lunch" dilemma [71] and suffer from the unidentifiability issue.
Corollary 3.6. (No Free Graph OOD Lunch) Without Assumption 3.3 or Assumption 3.5, there does not exist a learning algorithm that captures the invariance of the two-piece graph environments.
Corollary 3.6 is a natural conclusion from the previous discussion. The proof is straightforward and given in Appendix E.4. Assumption 3.3 and Assumption 3.5 establish the minimal premises for identifying the underlying invariant subgraphs. However, it also raises new challenges, as shown in Table.1. Chen et al. [8]