SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yongduo Sui $^{1\dagger}$ , Qitian Wu $^{2}$ , Jiancan Wu $^{1}$ , Qing Cui $^{3}$ , Longfei Li $^{3}$ , Jun Zhou $^{3*}$ , Xiang Wang $^{1*}$ , Xiangnan He $^{1*}$
<sup>1</sup>University of Science and Technology of China,
$^{2}$ Shanghai Jiao Tong University,  $^{3}$ Ant Group
syd2019@mail.ustc.edu.cn,echo740@sjtu.edu.cn
{wujcan, xiangwang1223, xiangnanhe}  $@$  gmail.com,
{cuiqing.cq,longyao.11f,jun.zhoujun}@antgroup.com
The issue of distribution shifts is emerging as a critical concern in graph representation learning. From the perspective of invariant learning and stable learning, a recently well-established paradigm for out-of-distribution generalization, stable features of the graph are assumed to causally determine labels, while environmental features tend to be unstable and can lead to the two primary types of distribution shifts. The correlation shift is often caused by the spurious correlation between environmental features and labels that differs between the training and test data; the covariate shift often stems from the presence of new environmental features in test data. However, most strategies, such as invariant learning or graph augmentation, typically struggle with limited training environments or perturbed stable features, thus exposing limitations in handling the problem of covariate shift. To address this challenge, we propose a simple-yet-effective data augmentation strategy, Adversarial Invariant Augmentation (AIA), to handle the covariate shift on graphs. Specifically, given the training data, AIA aims to extrapolate and generate new environments, while concurrently preserving the original stable features during the augmentation process. Such a design equips the graph classification model with an enhanced capability to identify stable features in new environments, thereby effectively tackling the covariate shift in data. Extensive experiments with in-depth empirical analysis demonstrate the superiority of our approach. The implementation codes are publicly available at https://github.com/yongduosui/AIA.
While recent advances have made solid progress in learning effective representations for graph-structured data, most of the existing approaches operate under the assumption that training and test graphs are independently drawn from an identical distribution [1, 2]. However, this assumption often falls short in real-world scenarios due to the out-of-distribution (OOD) data that potentially exists during the testing phase [3], which results in distribution shifts between training and test graphs. As a result, there is increasing research interest in OOD generalization on graphs or learning with distribution shifts on graphs [4]. Some of the typical recent works attempt to build effective methods for handling general distribution shifts on graphs, from (causal) invariant learning [5, 3], model architecture designs [6], and data augmentation [7].
This work was done during author's internship at Ant Group.
With more specific views, other attempts focus on designing generalizable models for tackling distribution shifts in particular domains with certain data formats, e.g., molecular graphs [8], recommender systems [9, 10], and anomaly detection [11].
However, the majority of existing studies primarily focus on the correlation shift, one type of distribution shift concerning OOD generalization [12, 13], leaving another equally important type of distribution shift, i.e., the covariate shift, largely under-explored in graph representation learning. From the perspective of invariant learning and stable learning [14, 15, 16], covariate shift is in stark contrast to correlation shift w.r.t. stable and environmental features of graph data. Specifically, according to the commonly-used graph generation hypothesis in prior studies [17, 18, 5, 8], there often exist stable features, which are informative features of the entire graphs and can reflect the predictive patterns in data. Based on this, the relationship between stable
features and labels is assumed to be invariant across environments. The remaining features could be unstable and varying across different environments, which mainly causes the following two distribution shifts: (1) correlation shift indicates that environments and labels establish inconsistent statistical correlations in training and test data, under the assumption that test environments are covered in the training data; whereas, (2) covariate shift characterizes that the environmental features in test data are unseen in training data [2, 12].
Considering a toy example in Figure 1, the environmental features ladder and tree are different in training and test data, which forms the covariate shift. Taking molecular property predictions as another example, functional groups (e.g., nitrogen dioxide  $(\mathrm{NO}_2)$ ) are stable to determine the predictive property of molecules [5, 8]. Whereas, scaffolds (e.g., carbon rings) are usually patterns irrelevant to the molecule properties, which can be seen as environmental features [2, 19]. In practice, we often need to use molecules collected in the past to train models, expecting that the models can predict the properties of molecules with new scaffolds in the future [20, 2, 19].
Considering the differences between correlation and covariate shifts, we take a close look into the existing efforts on graph generalization. They mainly fall into the following research lines, each of which has inherent limitations to solving the covariate shift. i) Invariant Graph Learning [17, 18, 5, 21] recently becomes a prevalent paradigm for OOD generalization. The basic idea is to capture stable features by minimizing the empirical risks in different environments. Unfortunately, it implicitly makes a prior assumption that test environments are available during training. This assumption is unrealistic owing to the obstacle of training data covering all possible test environments. Learning in limited environments can alleviate the spurious correlations that are hidden in the training data, but fail to extrapolate the test data with unseen environments. ii) Graph Data Augmentation [22, 23] perturbs graph features to enrich the distribution seen during training for better generalization. It can be roughly divided into node-level [24], edge-level [25], and graph-level [26, 7] with random [27] or adversarial strategies [28]. However, blindly augmenting the graphs can presumably destroys the stable features, and makes the augmented distributions out of control. For example, in Figure 1, the random strategy of DropEdge [25] will inevitably perturb the stable features (highlighted by red circles). As such, it may not sufficiently address the covariate shift and could potentially affect the generalization ability. Hence, we naturally ask a question: "Compared to the training data, can we generate new data that satisfy two conditions: 1) having new environments; 2) keeping the original stable features unchanged?"
Towards this end, we introduce two intuitive principles for graph augmentation: environmental feature discrepancy and stable feature consistency. The discrepancy principle promotes the exploration of new environments beyond the scope of training data, while the consistency principle seeks to maintain the integrity of stable features during augmentation. In order to achieve these principles, we devise a simple yet effective graph augmentation strategy: Adversarial Invariant Augmentation (AIA). Specifically, we employ an adversarial augmenter, a network that augments graphs by adversarially generating masks on them, thereby facilitating OOD exploration to enhance environmental discrep
Figure 1:  $P_{\text{train}}$  and  $P_{\text{test}}$  denote the training and test distributions.  $P_{\text{drop}}$  and  $P_{\text{ours}}$  represent the distributions of augmented data via DropEdge and AIA.
ancy. To foster stable feature consistency, we use another network, i.e., stable feature generator, which constructs masks that encapsulate stable features. We then delicately merge these masks and apply them to the graph data. As depicted in Figure 1, AIA primarily augments environmental features while leaving the stable elements unchanged. Our approach equips the graph classifier model with an enhanced ability to identify stable features in new environments and effectively mitigate the covariate shift issue. We also conduct extensive experiments and in-depth analyses. The experimental results highlight the limitations of several previous methods and underscore the superiority of our method in addressing the covariate shift issue, providing empirical support for our claims.
We define the uppercase letters  $(e.g., G)$  as random variables. The lower-case letters  $(e.g., g)$  are samples of variables, and the blackboard bold typefaces  $(e.g., \mathbb{G})$  denote the sample spaces. Let  $g = (\mathbf{A}, \mathbf{X}) \in \mathbb{G}$  denote a graph, where  $\mathbf{A}$  and  $\mathbf{X}$  are its adjacency matrix and node features, respectively. It is assigned with a label  $y \in \mathbb{Y}$  with a fixed labeling rule  $\mathbb{G} \to \mathbb{Y}$ . Let  $\mathcal{D} = \{(g_i, y_i)\}$  denote a dataset that is divided into a training set  $\mathcal{D}_{\mathrm{tr}} = \{(g_i^e, y_i^e)\}_{e \in \mathcal{E}_{\mathrm{tr}}}$  and a test set  $\mathcal{D}_{\mathrm{te}} = \{(g_i^e, y_i^e)\}_{e \in \mathcal{E}_{\mathrm{te}}}$ .  $\mathcal{E}_{\mathrm{tr}}$  and  $\mathcal{E}_{\mathrm{te}}$  are the index sets of training and test environments, respectively.
2.1 Definitions and Problem Formations
In this work, we focus on the graph classification scenario, which aims to train models with  $\mathcal{D}_{\mathrm{tr}}$  and infer the labels in  $\mathcal{D}_{\mathrm{te}}$ , such as molecular property prediction. From the viewpoints of invariant learning and stable learning, the inner mechanism of the labeling rule  $\mathbb{G} \to \mathbb{Y}$  is usually assumed to depend on the stable features [17, 18, 5, 8, 21], which are informative substructures of the entire graph. The relationship between the stable features and labels is assumed to be invariant across different environments, which makes OOD generalization possible [12]. Environmental features in graph data are assumed to have no causal-effect on labels. For instance, the chemical properties of molecules are mainly determined by specific functional groups, which can be regarded as stable features [17, 5, 19, 8]. Conversely, their scaffold structures, often irrelevant to their properties, can be seen as environmental features [20, 8].
Due to the instability of environmental features and the limitations in the data collection process, the training and test distributions are often inconsistent in real-world scenarios, i.e.,  $P_{\mathrm{tr}}(G,Y) \neq P_{\mathrm{te}}(G,Y)$ , which leads to two main types of distribution shifts [2, 12]: (1) Correlation shift (aka. spurious correlation or concept shift [2]) refers to  $P_{\mathrm{tr}}(G|Y) \neq P_{\mathrm{te}}(G|Y)$ ,  $P_{\mathrm{tr}}(G) = P_{\mathrm{te}}(G)$ . It indicates that there exist spurious statistical correlations in the training data, while these correlations might not hold in the test data. (2) Covariate shift denotes  $P_{\mathrm{tr}}(G|Y) = P_{\mathrm{te}}(G|Y)$ ,  $P_{\mathrm{tr}}(G) \neq P_{\mathrm{te}}(G)$ , which means that there exist new features, e.g., environmental features, in the test data. It may be ascribed to either insufficient quantity or diversity of data in the training set, as well as the unknown characteristics of the test environments. We provide formal definitions and examples of these two distribution shifts in Appendix A. Here, inspired by the prior study [12], we offer a formal definition to measure the graph covariate shift.
Definition 2.1 (Graph Covariate Shift) Let  $P_{\mathrm{tr}}$  and  $P_{\mathrm{te}}$  denote the probability functions of the training and test distributions. We measure the covariate shift between distributions  $P_{\mathrm{tr}}$  and  $P_{\mathrm{te}}$  as
$$
\operatorname {G C S} \left(P _ {\mathrm {t r}}, P _ {\mathrm {t e}}\right) = \frac {1}{2} \int_ {S} \left| P _ {\mathrm {t r}} (g) - P _ {\mathrm {t e}} (g) \right| d g, \tag {1}
$$
where  $S = \{g \in \mathbb{G} \mid P_{\mathrm{tr}}(g) \cdot P_{\mathrm{te}}(g) = 0\}$ , which covers the features (e.g., environmental features) that do not overlap between the two distributions.
$\mathrm{GCS}(\cdot ,\cdot)$  is always bounded in [0, 1]. The issue of graph covariate shift is very common in practice. For example, we often need to train models on past molecular graphs, and hope that the model can predict the chemical properties of future molecules with new features, e.g., new scaffolds [20]. In addressing the graph OOD issue, a majority of the strategies [17, 5, 18, 29, 21] grounded in invariant graph learning aim to pinpoint stable or invariant features. This is mainly accomplished by minimizing empirical risks across an array of training environments. Nonetheless, these methodologies frequently operate under the assumption of a shared input space across training and test data. This assumption, however, often fails on covariate shift. It presents substantial challenges for these models when it
comes to accurately identifying stable features within new testing environments. Consequently, while these methods typically exhibit satisfactory performance in managing correlation shifts, they often underperform in the face of covariate shifts. In this work, we focus on the covariate shift issue in graph classification, and we also give a formal definition of this problem as follows.
Problem 2.2 (Graph Classification under Covariate Shift) Given the training and test sets with environment sets  $\mathcal{E}_{\mathrm{tr}}$  and  $\mathcal{E}_{\mathrm{te}}$ , they follow distributions  $P_{\mathrm{tr}}$  and  $P_{\mathrm{te}}$ , and satisfy:  $\mathrm{GCS}(P_{\mathrm{tr}}, P_{\mathrm{te}}) > \epsilon$ , where  $\epsilon \in (0,1)$  represents the degree of covariate shift. We aim to use the data collected from training environments  $\mathcal{E}_{\mathrm{tr}}$ , and learn a powerful graph classifier  $f^{*}: \mathbb{G} \to \mathbb{Y}$  that performs well in all possible test environments  $\mathcal{E}_{\mathrm{te}}$ :
$$
f ^ {*} = \underset {f} {\arg \min } \sup  _ {e \in \mathcal {E} _ {\mathrm {t e}}} \mathbb {E} ^ {e} [ \ell (f (g), y) ], \tag {2}
$$
where  $\mathbb{E}^e [\ell (f(g),y)]$  is the empirical risk on the environment  $e$ , and  $\ell (\cdot ,\cdot)$  is the loss function.
To solve Problem 2.2, our idea is to generate new graphs through data augmentation. In this section, we first propose two principles for graph augmentation. Guided by these principles, we design a novel graph augmentation method, AIA, which can effectively address the covariate shift issue.
3.1 Two Principles for Graph Augmentation
We can observe that covariate shift is mainly caused by the scarcity of training environments. Hence, we first propose the discrepancy principle for graph augmentation.
Principle 3.1 (Environmental Feature Discrepancy) Given a graph set  $\{g\}$  with distribution function  $P$ , let  $T(\cdot)$  denote an augmentation function that augments graphs  $\{T(g)\}$  to distribution  $\widetilde{P}$ . Then  $T(\cdot)$  should meet  $\mathrm{GCS}(P, \widetilde{P}) \to 1$ .
From the perspective of data distribution, it requires that  $\widetilde{P}$  should keep away from the original distribution  $P$ . From the perspective of data instances, it emphasizes the discrepancy in the environments between the generated graphs and the original graphs. Since it does not give constraints on stable features, we here propose the second principle for graph augmentation.
**Principle 3.2 (Stable Feature Consistency)** Given a set of graphs  $\{g\}$  with a corresponding stable feature set  $\{g_{\mathrm{sta}} = (\mathbf{A}_{\mathrm{sta}}, \mathbf{X}_{\mathrm{sta}})\}$ . Let  $T(\cdot)$  denote an augmentation function that augments graphs  $\{T(g)\}$  with a corresponding stable feature set  $\{\widetilde{g}_{\mathrm{sta}} = (\widetilde{\mathbf{A}}_{\mathrm{sta}}, \widetilde{\mathbf{X}}_{\mathrm{sta}})\}$ . Then  $T(\cdot)$  should meet  $\mathbb{E}[\| \mathbf{A}_{\mathrm{sta}} - \widetilde{\mathbf{A}}_{\mathrm{sta}} \|_F^2] \to 0$  and  $\mathbb{E}[\| \mathbf{X}_{\mathrm{sta}} - \widetilde{\mathbf{X}}_{\mathrm{sta}} \|_F^2] \to 0$ , where  $\| \cdot \|_F$  is the Frobenius norm.
It necessitates that the stable features of the generated graphs should maintain consistency with those of the original graphs. This principle ensures the preservation of these stable features within the original training data, thereby safeguarding sufficient information pertaining to the labels. Consequently, this principle enhances the potential for generalization.
3.2 Out-of-distribution Exploration
Given a GNN model  $f(\cdot)$  with parameters  $\theta$ , we decompose  $f = \Phi \circ h$ , where  $h(\cdot): \mathbb{G} \to \mathbb{R}^d$  is a graph encoder to yield  $d$ -dimensional representations, and  $\Phi(\cdot): \mathbb{R}^d \to \mathbb{Y}$  is a classifier. To comply with Principle 3.1, we need to do OOD exploration. Inspired by distributionally robust optimization [30, 31], we consider the following optimization objective:
$$
\left. \min  _ {\theta} \left\{\sup  _ {\widetilde {P}} \left\{\mathbb {E} _ {\widetilde {P}} [ \ell (f (g), y) ]: D (\widetilde {P}, P) \leq \rho \right\} \right\}, \right. \tag {3}
$$
where  $P$  and  $\widetilde{P}$  are the original and explored data distributions, respectively.  $D(\cdot ,\cdot)$  is a distance metric between two probability distributions. The solution to Equation (3) guarantees the generalization within a robust radius  $\rho$  of the distribution  $P$ . To better measure the distance between distributions,
Figure 2: The overview of Adversarial Invariant Augmentation (AIA) Framework.
as suggested by [32], we adopt the Wasserstein distance [33, 34] as the distance metric. The distance metric function can be defined as  $D(\widetilde{P}, P) = \inf_{\mu \in \Gamma(\widetilde{P}, P)} \mathbb{E}_{\mu}[c(\widetilde{g}, g)]$ , where  $\Gamma(\widetilde{P}, P)$  is the set of all couplings of  $\widetilde{P}$  and  $P$ ;  $c(\cdot, \cdot)$  is the cost function. Studies [35, 34] also suggest that the distances in representation space typically correspond to semantic distances. Hence, we define the cost function in the representation space and give the transportation cost as  $c(\widetilde{g}, g) = \| h(\widetilde{g}) - h(g) \|_2^2$ . It denotes the "cost" of augmenting the graph  $g$  to  $\widetilde{g}$ . We can observe that it is difficult to set a proper  $\rho$ . Instead, we consider the Lagrangian relaxation for a fixed penalty coefficient  $\gamma$ . Inspired by [32], we can reformulate Equation (3) as follows:
$$
\min  _ {\theta} \left\{\sup  _ {\widetilde {P}} \left\{\mathbb {E} _ {\widetilde {P}} [ \ell (f (g), y) ] - \gamma D (\widetilde {P}, P) \right\} = \mathbb {E} _ {P} [ \phi (f (g), y) ] \right\}, \tag {4}
$$
where  $\phi(f(g), y) \coloneqq \sup_{\widetilde{g} \in \mathbb{G}} \{\ell(f(\widetilde{g}), y) - \gamma c(\widetilde{g}, g)\}$ . And we define  $\phi(f(g), y)$  as the robust surrogate loss. If we conduct gradient descent on the robust surrogate loss, we will have:
$$
\nabla_ {\theta} \phi (f (g), y) = \nabla_ {\theta} \ell (f (\widetilde {g} ^ {*}), y), \text {w h e r e} \widetilde {g} ^ {*} = \underset {\widetilde {g} \in \mathbb {G}} {\arg \max } \left\{\ell (f (\widetilde {g}), y) - \gamma c (\widetilde {g}, g) \right\}. \tag {5}
$$
$\widetilde{g}^*$  is an augmented view of the original data  $g$ . Hence, to achieve OOD exploration, we need to perform graph data augmentation via Equation (5) on the original data  $g$ .
3.3 Implementations of AIA
Equation (5) endows the ability of OOD exploration to data augmentation, which makes the augmented data meet the discrepancy principle. To achieve the consistency principle, we also need to capture stable features. Hence, we design a graph augmentation strategy: Adversarial Invariant Augmentation (AIA). The overview of the proposed framework is depicted in Figure 2, which mainly consists of two components: adversarial augmenter and stable feature generator. Adversarial augmenter achieves OOD exploration through adversarial data augmentation; meanwhile, the stable feature generator keeps stable feature consistency by identifying stable features from data. Below we elaborate on the implementation details.
Adversarial Augmenter & Stable Feature Generator. We design two networks, adversarial aug-. menter  $T_{\theta_1}(\cdot)$  and stable feature generator  $T_{\theta_2}(\cdot)$ , which generate masks for nodes and edges of graphs. They have the same structure and are parameterized by  $\theta_{1}$  and  $\theta_{2}$ , respectively. Given an input graph  $g = (\mathbf{A},\mathbf{X})$  with  $n$  nodes, mask generation network first obtains the node representations via a GNN encoder  $\widetilde{h} (\cdot)$ . To judge the importance of nodes and edges, it adopts two MLP networks  $\mathrm{MLP}_1(\cdot)$  and  $\mathrm{MLP}_2(\cdot)$  to generate the soft node mask matrix  $\mathbf{M}^x\in \mathbb{R}^{n\times 1}$  and edge mask matrix  $\mathbf{M}^a\in \mathbb{R}^{n\times n}$  for graph data, respectively. In summary, the mask generation network can be decomposed as:
$$
\mathbf {Z} = \widetilde {h} (g), \quad \mathbf {M} _ {i} ^ {x} = \sigma \left(\operatorname {M L P} _ {1} \left(\mathbf {h} _ {i}\right)\right), \quad \mathbf {M} _ {i j} ^ {a} = \sigma \left(\operatorname {M L P} _ {2} \left(\left[ \mathbf {z} _ {i}, \mathbf {z} _ {j} \right]\right)\right), \tag {6}
$$
where  $\mathbf{Z} \in \mathbb{R}^{n \times d}$  is node representation matrix, whose  $i$ -th row  $\mathbf{z}_i = \mathbf{Z}[i,:]$  denotes the representation of node  $i$ , and  $\sigma(\cdot)$  is the sigmoid function that maps the mask values  $\mathbf{M}_i^x$  and  $\mathbf{M}_{ij}^a$  to  $[0,1]$ .
Adversarial Invariant Augmentation. To estimate  $\widetilde{g}^*$  in Equation (5), we define the following adversarial learning objective:
$$
\max  _ {\theta_ {1}} \left\{\mathcal {L} _ {\mathrm {a d v}} = \mathbb {E} _ {P _ {\mathrm {t r}}} \left[ \ell \left(f \left(T _ {\theta_ {1}} (g)\right), y\right) - \gamma c \left(T _ {\theta_ {1}} (g), g\right) \right] \right\}. \tag {7}
$$
Then we can augment the graph by  $T_{\theta_1}(g) = (\mathbf{A} \odot \mathbf{M}_{\mathrm{adv}}^a, \mathbf{X} \odot \mathbf{M}_{\mathrm{adv}}^x)$ , where  $\odot$  is the broadcasted element-wise product. Although adversarially augmented graphs guarantee environmental discrepancy, it might destroys the stable parts. Therefore, we utilize the stable feature generator  $T_{\theta_2}(\cdot)$  to capture stable features and combine them with different environmental features. Following the sufficiency and invariance conditions [17, 3, 18, 5, 8, 29], we define the stable feature learning objective as:
$$
\min  _ {\theta , \theta_ {2}} \left\{\mathcal {L} _ {\mathrm {s t a}} = \mathbb {E} _ {P _ {\mathrm {t r}}} \left[ \ell \left(f \left(T _ {\theta_ {2}} (g)\right), y\right) + \ell \left(f (\widetilde {g}), y\right) \right] \right\}, \tag {8}
$$