SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shurui Gui, Meng Liu, Xiner Li, Youzhi Luo, Shuiwang Ji
{shurui.gui,mengliu,lxe,yzluo,sji}@tamu.edu
We tackle the problem of graph out-of-distribution (OOD) generalization. Existing graph OOD algorithms either rely on restricted assumptions or fail to exploit environment information in training data. In this work, we propose to simultaneously incorporate label and environment causal independence (LECI) to fully make use of label and environment information, thereby addressing the challenges faced by prior methods on identifying causal and invariant subgraphs. We further develop an adversarial training strategy to jointly optimize these two properties for causal subgraph discovery with theoretical guarantees. Extensive experiments and analysis show that LECI significantly outperforms prior methods on both synthetic and real-world datasets, establishing LECI as a practical and effective solution for graph OOD generalization. Our code is available at https://github.com/divelab/LECI.
Graph learning methods have been increasingly used for many diverse tasks, such as drug discovery [1], social network analysis [2], and physics simulation [3]. One of the major challenges in graph learning is out-of-distribution (OOD) generalization, where the training and test data are typically from different distributions, resulting in significant performance degradation when deploying models to new environments. Despite several recent efforts [4-7] have been made to tackle the OOD problem on graph tasks, their performances are often not satisfactory due to the complicated distribution shifts on graph topology and the violations of their assumptions or premises in reality. In addition, while environment-based methods have shown success in traditional OOD fields [8, 9], they cannot perform favorably on graphs [10]. Several existing graph environment-centered methods [11-13] aim to infer environment labels for graph tasks, however, most methods utilize the pre-collected environment labels in a manner similar to traditional environment-based techniques [9, 14, 15], rather than exploiting them in a unique and specific way tailored for graphs. Therefore, the potential exploitation of environment information in graph OOD tasks remains limited (Sec. 2.2). Although annotating or extracting environment labels requires additional cost, it has been proved that generalization without extra information is theoretically impossible [16] (Appx. B.2). This is analogous to the impossibility of inferring or approximating causal effects without counterfactual/intervened distributions [17, 18].
In this paper, we propose a novel learning strategy to incorporate label and environment causal independence (LECI) for tackling the graph OOD problem. Enforcing such independence properties can help exploit environment information and alleviate the challenging issue of graph topology shifts. Specifically, our contributions are summarized below. (1) We identify the current causal subgraph discovery challenges and propose to solve them by introducing label and environment causal independence. We further present a practical solution, an adversarial learning strategy, to
jointly optimize such two causal independence properties with causal subgraph discovery guarantees. (2) LECI is positioned as the first graph-specific pre-collected environment exploitation learning strategy. This learning strategy is applicable to any subgraph discovery networks and environment inference methods. (3) According to our extensive experiments, LECI outperforms baselines on both the structure/feature shift sanity checks and real-world scenario comparisons. With additional visualization, hyperparameter sensitivity, training dynamics, and ablation studies, LECI is empirically proven to be a practically effective method. (4) Contrary to prevalent beliefs and results, we showcase that pre-collected environment information, far from being useless, can be a potent tool in graph tasks, which is evidenced more powerful than previous assumptions.
2.1 Graph OOD generalization
We represent an attributed graph as  $G = (X, A) \in \mathcal{G}$ , where  $\mathcal{G}$  is the graph space.  $X \in \mathbb{R}^{n \times d}$  and  $A \in \mathbb{R}^{n \times n}$  denote its node feature matrix and adjacency matrix respectively, where  $n$  is the number of nodes and  $d$  is the feature dimension. In the graph-level OOD generalization setting, each graph  $G$  has an associated label  $Y \in \mathcal{V}$ , where  $\mathcal{V}$  denotes the label space. Notably, training and test data are typically drawn from different distributions, i.e.,  $P^{tr}(G, Y) \neq P^{te}(G, Y)$ . Unlike the image space, where distribution shifts occur only on feature vectors, distribution shifts in the graph space can happen on more complicated variables such as graph topology [19, 20]. Therefore, many existing graph OOD methods [4-7, 11] aim to identify the most important topological information, called causal subgraphs, so that they can be used to make predictions that are robust to distribution shifts.
2.2 Environment-based OOD algorithms
To address distribution shifts, following invariant causal predictor (ICP) [21] and invariant risk minimization (IRM) [9], many environment-based invariant learning algorithms [14, 22-24], also referred as parts of domain generalization (DG) [25-27], classify data into several groups called environments, denoted as  $E \in \mathcal{E} = \{e_1, e_2, \dots, e_{|\mathcal{E}|}\}$ . Intuitively, data in the same environment group share similar uncritical information, such as the background of images and the size of graphs. It is assumed that the shift between the training and test distribution should be reflected among environments. Therefore, to address generalization problems in the test environments, these methods incorporate the environment shift information from training environments into neural networks, thereby driving the networks to be invariant to the shift between the training and test distribution.
Several works further consider the inaccessibility of environment information in OOD settings. For example, environment inference methods [11, 12, 28] propose to infer environment labels to make invariant predictions. While in graph tasks, several methods [4-6, 13] attempt to generalize without the use of environments. However, these algorithms usually rely on relatively strict assumptions that can be compromised in many real-world scenarios and may be more difficult to satisfy than the access of environment information, which will be further compared and justified in Appx C.2. As demonstrated in the Graph Out-of-Distribution (GOOD) benchmark [10] and DrugOOD [29], the environment information for graph datasets is commonly accessible. Thus, in this paper, we assume the availability of environment labels and focus on invariant predictions by exploiting the given environment information in graph-specific OOD settings. Extensive discussions of related works and comparisons to previous graph OOD methods are available in Appx B.
Comparisons to other environment-based graph OOD methods. It's crucial to distinguish between the two stages of environmental-based methods: environmental inference and environmental exploitation. The first phase, environmental inference, involves predicting environmental labels, while the second, environmental exploitation, focuses on using pre-acquired environmental labels. Typically, an environmental inference method employs an environmental exploitation method to evaluate its effectiveness. However, an environmental exploitation method does not require an environmental inference method. Recent developments in graph-level environmental inference methods [11, 12] and node-level environmental inference methods [13] introduce graph-specific environmental inference techniques. Nevertheless, their corresponding environmental exploitation strategies are not tailored to graphs. In contrast, our method, LECI, bypasses the environmental inference phase and instead introduces a graph-specific environmental exploitation algorithm, supported by justifications in Appx B.3. More environment-related discussions and motivations are available in Appx. B.2.
OOD generalization is a longstanding problem because different distribution shifts exist in many different applications of machine learning. As pointed out by Kaur et al. [30], understanding the data generation process is critical for solving OOD problems. Hence, we first formalize the target problems through three topology distribution shift assumptions from a causal and data generation perspective. We further identify the challenges of discovering causal subgraphs for addressing topology distribution shifts. Building on our analysis, we propose a technical solution and a practical implementation to jointly optimize label and environment causal independence in order to learn causal subgraphs effectively.
3.1 Causal perspective of graph OOD generalization
A major difference between traditional OOD and graph OOD tasks is the topological structure distribution shift in graph OOD tasks. To analyze graph structure shifts, we commonly assume that only a part of a graph determines its target  $Y$ ; e.g., only the functional motif of a molecule determines its corresponding property. Therefore, in graph distribution shift studies, we assume that each graph  $G \in \mathcal{G}$  is composed of two subgraphs; namely, a causal subgraph  $G_{C} \in \mathcal{G}$  and a spurious subgraph  $G_{S} \in \mathcal{G}$  as shown in the structure causal models (SCMs) [17, 18] of Fig. 1. For notational convenience, we denote graph union and subtraction operations as  $+$  and  $-$ , respectively; e.g.
(a) Covariate SCM
(c) PIIF SCM
Figure 1: Illustrations of three distribution shift assumptions. In these structural causal models (SCMs), each node denotes a variable, and each directional edge represents causation. Grey nodes are observable variables that can be directly accessed, but it does not mean that we are "conditioned on" these variables.
$G = G_{C} + G_{S}$  and  $G_{C} = G - G_{S}$ . It follows from the above discussions that  $G_{C}$  is the only factor determining the target variable  $Y \in \mathcal{V}$ , while  $G_{S}$  is controlled by an exogenous environment  $E$ , such that graphs in the same environment share similar spurious subgraphs.
As illustrated in Fig. 1 (a), the covariate shift represents the most common distribution shift. Similar to image problems in which the same object is present with different backgrounds, graphs that share the same functional motif along with diverse graph backbones belong to the covariate shift. Under this setting,  $G_{S}$  is only controlled by  $E$ . Models trained in this setting suffer from distribution shifts because the general empirical risk minimization (ERM) training objective  $P(Y|G)$  is conditioned on the collider  $G$ , which builds a spurious correlation between  $G_{S}$  and  $Y$  through  $Y \gets G_{C} \to G \gets G_{S}$ , where  $\rightarrow$  denotes a causal correlation. The fully informative invariant features (FIIF) and the partially informative invariant features (PIIF) assumptions, shown in Fig. 1 (b) and (c), are two other common assumptions proposed by invariant learning studies [6, 9, 31]. In the graph FIIF assumption,  $G_{S}$  is controlled by both  $G_{C}$  and  $E$ , which constructs an extra spurious correlation  $Y \gets G_{C} \to G_{S}$ . In contrast, the graph PIIF assumption introduces an anti-causal correlation between  $G_{S}$  and  $Y$ . Our work focuses on the covariate shift generalization to develop our approach and then extends the solution to both FIIF and PIIF assumptions.
Essentially, we are addressing the distribution shifts between  $P^{tr}(G,Y)$  and  $P^{te}(G,Y)$ . Our core assumption, based on the Independence Causal Mechanism (ICM), is that there exists a component  $G_{C}$  within  $G$  such that  $P^{tr}(Y|G_C) = P^{te}(Y|G_C)$ . The shifts in distribution are exclusively attributed to interventions on  $G_{S} = G - G_{C}$  that are encapsulated as an environment variable  $E$ , aligned with the invariant learning literature. The OOD dilemma arises when  $P^{tr}(E) \neq P^{te}(E)$ , leading to shifts in both  $P^{tr}(G_S)$  and  $P^{te}(G_S)$ , and consequentially,  $P^{tr}(G) \neq P^{te}(G)$ . It's crucial to note that we assume both  $P^{tr}(G_C)$  and  $P^{te}(G_C)$  retain the same support. Explicitly, our target is to resolve OOD scenarios where the supports of  $P^{tr}(E)$  and  $P^{te}(E)$  differ.
3.2 Subgraph discovery challenges
A common strategy to make neural networks generalizable to structure distribution shifts in graph OOD tasks is to identify causal subgraphs for invariant predictions [4-6, 19]. However, correctly selecting subgraphs is challenging due to the following two precision challenges. (1) The selected causal subgraph may contain spurious structures, i.e.,  $\exists G_{p} \subseteq G_{S}, G_{p} \subseteq \hat{G}_{C}$ , where  $G_{a} \subseteq G_{b}$
Figure 2: An illustration of the invariant prediction process by selecting the causal subgraph  $\hat{G}_C$  using an interpretable subgraph discovery network in LECI. This network and training details are described in Sec. 3.4. We multiply the reversed gradients with two hyperparameters  $\lambda_{\mathrm{L}}$  and  $\lambda_{\mathrm{E}}$  to control the adversarial training. For clear notations, we define the three losses  $\mathcal{L}_{\mathrm{inv}}$ ,  $\mathcal{L}_{\mathrm{L}}$ , and  $\mathcal{L}_{\mathrm{E}}$  in Eq. 1, 7, and 6, respectively.
denotes  $G_{a}$  is a subgraph of  $G_{b}$ , and  $\hat{G}_{C}$  represents the selected causal subgraph. (2) The model may not select the whole  $G_{C}$  into  $\hat{G}_{C}$  and leave a part of  $G_{C}$  in  $\hat{G}_{S} = G - \hat{G}_{C}$ . Formally,  $\exists G_{p} \subseteq G_{C}, G_{p} \subseteq \hat{G}_{S}$ . The reason for the occurrence of these precision challenges will be further discussed in Appx. C.3.
3.3 Two causal independence properties
To address the precision issues, it is necessary to distinguish between causal subgraphs and spurious subgraphs, both theoretically and in practice. To achieve this, under the covariate assumption, we introduce two causal independence properties for causal and spurious subgraphs, respectively; those are  $E \perp G_C$  and  $Y \perp G_S$ .
The first property  $E \perp G_C$  is a crucial consideration to alleviate the first precision problem. As illustrated in Fig. 1 (a), since  $G$  acts as a collider that blocks the correlation between  $E$  and  $G_C$ , the environment factor  $E$  should be independent of the causal subgraph  $G_C$ . Conversely, due to the direct causation between  $E$  and  $G_S$ , the spurious subgraph  $G_S$  is highly correlated with  $E$ . This correlation difference with  $E$  indicates that enforcing independence between  $E$  and the selected causal subgraph  $\hat{G}_C$  can prevent parts of spurious subgraphs in  $\hat{G}_C$  be included, thereby solving the first precision problem. Note that this independence also holds in FIIF and PIIF settings, because  $G_S$  acts as another collider and blocks the left correlation paths between  $G_C$  and  $E$ , i.e.,  $G_C \rightarrow G_S \leftarrow E$  and  $G_C \rightarrow Y \rightarrow G_S \leftarrow E$ .
We propose to enforce the second independence property  $Y \perp G_S$  in order to address the second precision problem. Under the covariate assumption, the correlation between  $Y$  and  $G_S$  is blocked by  $G$ , leading to the independence between  $Y$  and  $G_S$ . On the other side, the causal subgraph  $G_C$  is intrinsically correlated with  $Y$ . Similarly, this correlation difference with  $Y$  motivates us to enforce such independence, thus filtering parts of causal subgraphs out of the selected spurious subgraphs  $\hat{G}_S$ . This property  $Y \perp G_S$ , however, does not hold under the FIIF and PIIF settings, because of the  $Y \gets G_C \to G_S$  and the  $Y \to G_S$  correlations. Therefore, we introduce a relaxed version of the independence property: for any  $G_p$  that  $G_S \subseteq G_p$ ,  $I(G_S;Y) \leq I(G_p;Y)$  (See Appx. C.1.1).
3.4 Adversarial implementation
The above two causal independence properties are intuitively helpful in discovering causal subgraphs. In this section, we present our practical solution to learn the label and environment causal independence (LECI). The overall architecture of LECI is illustrated in Fig. 2. To be specific, we use an interpretable subgraph discovery network as the basic architecture, which consists of a subgraph selector  $f_{\theta} : \mathcal{G} \mapsto \mathcal{G}$  and an invariant predictor  $g_{\mathrm{inv}} : \mathcal{G} \mapsto \mathcal{V}$ . Here,  $f_{\theta}$  and  $g_{\mathrm{inv}}$  model the distributions  $P_{\theta}(G_C|G)$  and  $P_{\phi_{\mathrm{inv}}}(\hat{Y}|\hat{G}_C)$ , where  $\phi_{\mathrm{inv}}$  is the model parameters. Specifically, following Miao et al.
[5], our subgraph selector module selects subgraphs by sampling edges from the original graphs. However, the sampling process is not differentiable and blocks the gradient back-propagation. To overcome this problem, we use Gumbel-Sigmoid [32] to bypass the non-differentiable sampling process. The basic training loss for this subgraph discovery network can be written as
$$
\mathcal {L} _ {\text {i n v}} = - \mathbb {E} \left[ \log P _ {\phi_ {\text {i n v}}} (Y | \hat {G} _ {C}) \right]. \tag {1}
$$
It is worth noting that  $\hat{G}_C = f_\theta(G)$  and  $\hat{G}_S = G - \hat{G}_C$  are complimentary. More details of the subgraph selector can be found in Appx. D.1
To incorporate the causal independence properties described in Sec. 3.3, we first enforce the condition  $\hat{G}_C \perp E$ . Since  $\hat{G}_C \perp E$  is equivalent to  $I(E; \hat{G}_C) = 0$ , and  $I(E; \hat{G}_C) \geq 0$ , the objective  $I(E; \hat{G}_C)$  reaches its minimum when and only when  $\hat{G}_C \perp E$ . This leads us to define minimizing  $I(E; \hat{G}_C)$  as our training criterion.
Definition 3.1. The environment independence training criterion is
$$
\theta_ {\mathrm {E}} ^ {*} = \underset {\theta} {\arg \min } I (E; \hat {G} _ {C}). \tag {2}
$$
For a fixed  $\theta$ ,  $I(E; \hat{G}_C)$  is essentially a constant. However, the mutual information  $I(E; \hat{G}_C) = \mathbb{E}\left[\log \frac{P(E|\hat{G}_C)}{P(E)}\right]$  cannot be calculated directly, since  $P(E|\hat{G}_C)$  is unknown. Following Proposition 1 in GAN [33], we introduce an optimal discriminator  $g_{\mathrm{E}}: \mathcal{G} \mapsto \mathcal{E}$  with parameters  $\phi_E$  to approximate  $P(E|\hat{G}_C)$  as  $P_{\phi_{\mathrm{E}}} (E|\hat{G}_C)$ , minimizing the negative log-likelihood  $-\mathbb{E}\left[\log P_{\phi_{\mathrm{E}}} (E|\hat{G}_C)\right]$ . We then have the following two propositions:
Proposition 3.2. For  $\theta$  fixed, the optimal discriminator  $\phi_E$  is
$$
\phi_ {E} ^ {*} = \underset {\phi_ {E}} {\arg \min } - \mathbb {E} \left[ \log P _ {\phi_ {E}} (E | \hat {G} _ {C}) \right]. \tag {3}
$$
Proposition 3.3. Denoting  $KL$ -divergence as  $KL[\cdot \| \cdot ]$ , for  $\theta$  fixed, the optimal discriminator  $\phi_E$  is  $\phi_{E}^{*}$ , s.t.
$$
K L \left[ P (E | \hat {G} _ {C}) \| P _ {\phi_ {E} ^ {*}} (E | \hat {G} _ {C}) \right] = 0. \tag {4}
$$
Proposition 3.2 can be proved straightforwardly by applying the cross-entropy training criterion, while the proof of Proposition 3.3 is provided in Appx C. With these propositions, the mutual information can be computed with the help of the optimal discriminator  $\phi_E^*$ . According to Proposition 3.3, we have:
$$
\begin{array}{l} I (E; \hat {G} _ {C}) = \mathbb {E} \left[ \log P _ {\phi_ {\mathrm {E}} ^ {*}} (E | \hat {G} _ {C}) \right] + H (E) + \mathbf {K L} \left[ P (E | \hat {G} _ {C}) \| P _ {\phi_ {\mathrm {E}} ^ {*}} (E | \hat {G} _ {C}) \right] \\ = \mathbb {E} \left[ \log P _ {\phi_ {\mathrm {E}} ^ {*}} (E | \hat {G} _ {C}) \right] + H (E) + 0. \tag {5} \\ \end{array}
$$
Thus, by disregarding the constant  $H(E)$ , the training criterion becomes:
$$
\theta_ {\mathrm {E}} ^ {*} = \underset {\theta} {\arg \min } I (E; \hat {G} _ {C}) = \underset {\theta} {\arg \min } \mathbb {E} \left[ \log P _ {\phi_ {\mathrm {E}} ^ {*}} (E | \hat {G} _ {C}) \right] = \underset {\theta} {\arg \min } \left\{\underset {\phi_ {E}} {\max } \mathbb {E} \left[ \log P _ {\phi_ {\mathrm {E}}} (E | \hat {G} _ {C}) \right] \right\}. \tag {6}
$$
We can reformulate this training criterion in terms of a negative log-likelihood form  $(\mathcal{L}_{\mathrm{E}})$  to define our environment adversarial training criterion (EA):
$$
\Theta_ {\mathrm {E}} ^ {*} = \left\{\theta_ {\mathrm {E}} ^ {*}: \theta_ {\mathrm {E}} ^ {*} \in \underset {\theta} {\arg \max } \left\{\underset {\phi_ {E}} {\min } \mathcal {L} _ {\mathrm {E}} \right\} \right\} = \left\{\theta_ {\mathrm {E}} ^ {*}: \theta_ {\mathrm {E}} ^ {*} \in \underset {\theta} {\arg \max } \left\{\underset {\phi_ {E}} {\min } - \mathbb {E} \left[ \log P _ {\phi_ {\mathrm {E}}} (E | \hat {G} _ {C}) \right] \right\} \right\}. \tag {7}
$$
To enforce  $\hat{G}_S \perp Y$ , we introduce a label discriminator  $g_{\mathrm{L}}: \mathcal{G} \mapsto \mathcal{V}$  to model  $P(Y|\hat{G}_S)$  with parameters  $\phi_{\mathrm{L}}$  as  $P_{\phi_{\mathrm{L}}} (Y|\hat{G}_{S})$ . This leads to a symmetric label adversarial training criterion (LA):
$$
\Theta_ {\mathrm {L}} ^ {*} = \left\{\theta_ {\mathrm {L}} ^ {*}: \theta_ {\mathrm {L}} ^ {*} \in \arg \max  _ {\theta} \left\{\min  _ {\phi_ {L}} \mathcal {L} _ {\mathrm {L}} \right\} \right\} = \left\{\theta_ {\mathrm {L}} ^ {*}: \theta_ {\mathrm {L}} ^ {*} \in \arg \max  _ {\theta} \left\{\min  _ {\phi_ {L}} - \mathbb {E} \left[ \log P _ {\phi_ {\mathrm {L}}} (Y | \hat {G} _ {S}) \right] \right\} \right\}. \tag {8}
$$
It's important to note that  $\Theta_{\mathrm{E}}^{*}$  and  $\Theta_{\mathrm{L}}^{*}$  are not unique optimal parameters. Instead, they are two optimal parameter sets, including parameters that satisfy the training criteria. We will discuss  $\Theta_{\mathrm{E}}^{*}$  and  $\Theta_{\mathrm{L}}^{*}$ , along with a relaxed version of LA, in further detail in the theoretical analysis section 3.5.
Intuitively, as illustrated in Fig. 2, the models  $g_{\mathrm{E}}$  and  $g_{\mathrm{L}}$ , acting as an environment discriminator and a label discriminator, are optimized to minimize the negative log-likelihoods. In contrast, the subgraph selector  $f_{\theta}$  tries to adversarially maximize the losses by reversing the gradients during backpropagation. Through this adversarial training process, we aim to select a causal subgraph  $\hat{G}_C$  that is independent of the environment variables  $E$ , while simultaneously eliminating causal information from  $\hat{G}_S$  such that the selected causal subgraph  $\hat{G}_C$  can retain as much causal information as possible for the accurate inference of the target variable  $Y$ .
3.5 Theoretical results
Under the framework of the three data generation assumptions outlined in Section 3.1, we provide theoretical guarantees to address the challenges associated with subgraph discovery. Detailed proofs for the following lemmas and theorems can be found in Appx. C. We initiate our analysis with a lemma for  $\hat{G}_C \perp E$ :
Lemma 3.4. Given a subgraph  $G_p$  of input graph  $G$  and SCMs (Fig. 1), it follows that  $G_p \subseteq G_C$  if and only if  $G_p \perp E$ .
This proof hinges on the data generation assumption that any substructures external to  $G_{C}$  (i.e.,  $G_{S}$ ) maintain associations with  $E$ . Consequently, any  $\hat{G}_{C}$  that fulfills  $\hat{G}_{C} \perp E$  will be a subgraph of  $G_{C}$ , denoted as  $\hat{G}_{C} \subseteq G_{C}$ . We then formulate the EA training lemma as follows:
Lemma 3.5. Given SCMs (Fig. 1), the predicted  $\hat{G}_C = f_\theta(G)$ , and the EA training criterion (Eq. 6),  $\hat{G}_C \subseteq G_C$  if and only if  $\theta \in \Theta_E^*$ .