Title: Unleashing the Power of Graph Data Augmentation on Covariate Distribution Shift
Venue: NeurIPS 2023
Accepted: yes
Keywords: Graph Neural Network

Graph Data Augmentation

Distribution Shift
Abstract: 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.

Unleashing the Power of Graph Data Augmentation on Covariate Distribution Shift

Abstract

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 .

1 Introduction

While recent advances have made solid progress in learning effective representations for graphstructured 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].

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

…

…

…

…

𝑃train

𝑃test

𝑃drop

𝑃 ours

Training

graphs

Test

graphs

DropEdge

graphs

AIA

graphs

Environmental features

(ladder & tree)

stable feature (house)

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 (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-

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.

2 Preliminaries

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., G ) denote the sample spaces. Let g = ( A , X ) ∈ G denote a graph, where A and X are its adjacency matrix and node features, respectively. It is assigned with a label y ∈ Y with a fixed labeling rule G → Y . Let D = {( g i , y i )} denote a dataset that is divided into a training set D tr = {( g e i , y e i )} e ∈E tr and a test set D te = {( g e i , y e i )} e ∈E te . E tr and E 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 D tr and infer the labels in D te , such as molecular property prediction. From the viewpoints of invariant learning and stable learning, the inner mechanism of the labeling rule G → 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 tr ( G,Y ) ≠ P 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 tr ( G ∣ Y ) ≠ P te ( G ∣ Y ) , P tr ( G ) = P 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 tr ( G ∣ Y ) = P te ( G ∣ Y ) , P tr ( G ) ≠ P 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 tr and P te denote the probability functions of the training and test distributions. We measure the covariate shift between distributions P tr and P te as

where S = { g ∈ G ∣ P tr ( g ) ⋅ P te ( g ) = 0 } , which covers the features (e.g., environmental features) that do not overlap between the two distributions.

GCS (⋅ , ⋅) 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 E tr and E te , they follow distributions P tr and P te , and satisfy: GCS ( P tr , P te ) > ϵ , where ϵ ∈ ( 0 , 1 ) represents the degree of covariate shift. We aim to use the data collected from training environments E tr , and learn a powerful graph classifier f ∗ ∶ G → Y that performs well in all possible test environments E te :

where E e [ ℓ ( f ( g ) , y )] is the empirical risk on the environment e , and ℓ (⋅ , ⋅) is the loss function.

3 Methodology

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 (⋅) denote an augmentation function that augments graphs { T ( g )} to distribution ̃ P . Then T (⋅) should meet GCS ( P, ̃ P ) → 1 .

From the perspective of data distribution, it requires that ̃ 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 sta = ( A sta , X sta )} . Let T (⋅) denote an augmentation function that augments graphs { T ( g )} with a corresponding stable feature set {̃ g sta = ( ̃ A sta , ̃ X sta ) } . Then T (⋅) should meet E [∥ A sta - ̃ A sta ∥ 2 F ] → 0 and E [∥ X sta - ̃ X sta ∥ 2 F ] → 0 , where ∥ ⋅ ∥ 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 (⋅) with parameters θ , we decompose f = Φ ○ h , where h (⋅) ∶ G → R d is a graph encoder to yield d -dimensional representations, and Φ (⋅) ∶ R d → 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:

where P and ̃ P are the original and explored data distributions, respectively. D (⋅ , ⋅) is a distance metric between two probability distributions. The solution to Equation (3) guarantees the generalization within a robust radius ρ of the distribution P . To better measure the distance between distributions,

Adversarial

Augmenter

Stable Feature

Generator

Augmented Graphs

Training time

· · ·

Stable Features

Environmental Features

GNN

Classifier

· · ·

· · ·

· · ·

as suggested by [32], we adopt the Wasserstein distance [33, 34] as the distance metric. The distance metric function can be defined as D ( ̃ P,P ) = inf µ ∈ Γ ( ̃ P,P ) E µ [ c (̃ g, g )] , where Γ ( ̃ P,P ) is the set of all couplings of ̃ P and P ; c (⋅ , ⋅) 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 (̃ g, g ) = ∥ h (̃ g ) -h ( g )∥ 2 2 . It denotes the 'cost' of augmenting the graph g to ̃ g . We can observe that it is difficult to set a proper ρ . Instead, we consider the Lagrangian relaxation for a fixed penalty coefficient γ . Inspired by [32], we can reformulate Equation (3) as follows:

where ϕ ( f ( g ) , y ) ∶= sup ̃ g ∈ G { ℓ ( f (̃ g ) , y ) -γc (̃ g, g )} . And we define ϕ ( f ( g ) , y ) as the robust surrogate loss. If we conduct gradient descent on the robust surrogate loss, we will have:

̃ 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 augmenter T θ 1 (⋅) and stable feature generator T θ 2 (⋅) , which generate masks for nodes and edges of graphs. They have the same structure and are parameterized by θ 1 and θ 2 , respectively. Given an input graph g = ( A , X ) with n nodes, mask generation network first obtains the node representations via a GNN encoder ̃ h (⋅) . To judge the importance of nodes and edges, it adopts two MLP networks MLP 1 (⋅) and MLP 2 (⋅) to generate the soft node mask matrix M x ∈ R n × 1 and edge mask matrix M a ∈ R n × n for graph data, respectively. In summary, the mask generation network can be decomposed as:

where Z ∈ R n × d is node representation matrix, whose i -th row z i = Z [ i, ∶] denotes the representation of node i , and σ (⋅) is the sigmoid function that maps the mask values M x i and M a ij to [ 0 , 1 ] .

Adversarial Invariant Augmentation. To estimate ̃ g ∗ in Equation (5), we define the following adversarial learning objective:

Then we can augments the graph by T θ 1 ( g ) = ( A ⊙ M a adv , X ⊙ M x adv ) , where ⊙ 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 θ 2 (⋅) 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:

where ̃ g = ( A ⊙ ̃ M a , X ⊙ ̃ M x ) is the augmented graph. It adopts the mask combination strategy: ̃ M a = ( 1 a -M a sta ) ⊙ M a adv + M a sta and ̃ M x = ( 1 x -M x sta ) ⊙ M x adv + M x sta , where M a sta and M x sta are generated by T θ 2 (⋅) , 1 a and 1 x are all-one matrices, and if there is no edge between node i and node j , then we set 1 a ij to 0. Now we explain this combination strategy. Taking ̃ M x as an example, since M x sta denotes the captured stable regions via T θ 2 (⋅) , 1 x -M x sta represents the complementary parts, which are environmental regions. M x adv represents the adversarial perturbation, so ( 1 x -M x sta ) ⊙ M x adv is equivalent to applying the adversarial perturbation on environmental features, meanwhile, sheltering the stable features. Finally, + M sta signifies that the augmented data should preserve the original stable features. Consequently, it satisfies both principles. Upon analysis of Equation (8), the first term implies that the stable features are sufficient for making right predictions. The second term promotes right and invariant predictions under generated environments utilizing stable features.

Regularization. For Equation (7), the adversarial optimization tends to remove more nodes and edges, so we should also constrain the perturbations. Although Equation (8) satisfies the sufficiency and invariance conditions, it is necessary to impose constraints on the ratio of the stable features to prevent trivial solutions. Hence, we first define the regularization function r ( M , k, λ ) = (∑ ij M ij / k -λ ) + (∑ ij I [ M ij > 0 ]/ k -λ ) , where k is the total number of elements to be constrained, I ∈ { 0 , 1 } is an indicator function. The first term penalizes the average ratio close to λ , while the second term encourages an uneven distribution. Given a graph with n nodes and m edges, we define the regularization term for adversarial augmentation and stable feature learning as:

where λ s ∈ ( 0 , 1 ) is the ratio of stable features, we usually set λ a = 1 for adversarial learning, which can alleviate excessive perturbations. The algorithm is provided in Appendix D.1.

4 Theoretical Discussions

In this section, we engage in theoretical discussions to elucidate our learning objective and its connections with the covariate shift. We first explore the relationship between our optimization objective and the discrepancy principle. Recalling the optimization objective of Equation 3, we aspire to identify a distribution ̃ P that can manifest within a Wasserstein ball [36], which is centered on distribution P , with distance ρ serving as the radius. Under appropriate conditions, we find that our learning optimization objective can establish a close connection with OOD exploration.

Proposition 4.1 Consider a probability distribution P defined over a measurable space ( Ω , F) , where Ω denotes the sample space and F is a σ -algebra on Ω . We construct two Wasserstein balls with P at their center and radii ρ 1 and ρ 2 respectively. Utilizing Equation 3, we generate two distinct distributions, ̃ P 1 and ̃ P 2 , within the space ( Ω , F) . If (i) P is an isotropic distribution; (ii) ∃ x 1 , x 2 ∈ Ω such that P ( x ) = ̃ P 1 ( x -x 1 ) = ̃ P 2 ( x -x 2 ) ; (iii) ρ 1 < ρ 2 , then we have GCS ( P, ̃ P 1 ) ≤ GCS ( P, ̃ P 2 ) .

This suggests that by appropriately increasing the robustness radius in AIA, we can effectively amplify the covariate shift between the training and generated distributions. This in turn underscores the reliability of our discrepancy principle, to a certain degree. Comprehensive proofs and detailed discussions supporting these conclusions can be found in Appendix B.

5 Experiments

In this section, we conduct extensive experiments to answer the following R esearch Q uestions:

RQ1: Compared to existing efforts, how does AIA perform under covariate shift?

RQ2: Can the proposed AIA achieve the principles of environmental feature discrepancy and stable feature consistency, thereby alleviating the graph covariate shift?

RQ3: How do the different components and hyperparameters of AIA affect the performance?

5.1 Experimental Settings

Datasets. We use graph OOD datasets [2] and OGB datasets [20], which include Motif, CMNIST, Molbbbp, and Molhiv. Following [2], we adopt the base, color, size, and scaffold data splitting to create various covariate shifts. The details of the datasets, metrics, implementations, and other settings are provided in Appendix D.2 and D.4. More experiments are provided in Appendix E.

Baselines. We adopt 16 baselines, which can be divided into the following three specific categories:

General Generalization Algorithms: ERM, IRM [14], GroupDRO [31], VREx [37].

Graph Generalization Algorithms: DIR [17], CAL [5], GSAT [38], OOD-GNN [39], StableGNN [40], CIGA [41], DisC [21].

Graph Augmentation: DropEdge [25], GREA [18], FLAG [24], M-Mixup [26], G -Mixup [7].

5.2 Main Results (RQ1)

We first make comparisons with various baselines in Table 1, and have the following observations:

Most generalization and augmentation methods fail under covariate shift. VREx achieves a 2.81% improvement on Motif (base). For two shifts of Molhiv, data augmentation methods GREA and DropEdge obtain 1.20% and 0.77% improvements. The invariant learning methods, i.e., DIR and CAL also obtain 4.60% and 1.53% improvements on CMNIST and Molbbbp (size). Unfortunately, none of the methods consistently outperform ERM. For example, GREA and DropEdge perform poorly on Motif (base), ↓ 11.92% and ↓ 23.58%. DIR and CAL also fail on Molhiv. These show that both invariant learning and data augmentation methods have their own weaknesses, which lead to unstable performance when facing complex and diverse covariate shifts from different datasets.

AIA consistently outperforms most baseline methods. Compared with ERM, AIA can obtain significant improvements. For two types of covariate shifts on Motif, AIA surpasses ERM by 4.98% and 4.11%, respectively. In contrast to the large performance variances on different datasets achieved by

Original graphs

…

…

Augmented graphs

baselines, AIA consistently obtains the leading performance across the board. For CMNIST, AIA achieves a performance improvement of 3.17% compared to the best baseline DIR. For Motif, the performance is improved by 2.17% and 1.72% compared to VREx and GREA. These results illustrate that AIA can overcome the shortcomings of invariant learning and data augmentation. Armed with the principles of environmental feature diversity and stable feature invariance, AIA achieves stable and consistent improvements on different datasets with various covariate shifts. In addition, although we focus on covariate shift in this work, we also carefully check the performance of AIA under correlation shift, and the results are presented in Appendix E.

5.3 In-depth Analyses (RQ2)

In this section, we conduct qualitative and quantitative experiments to support our two principles. Firstly, we utilize GCS (⋅ , ⋅) as the measurement to quantify the degree of covariate shift. The detailed estimation procedure is provided in Appendix C. We select four different domains, i.e., base, size, color and scaffold, to create covariate shifts. The experimental results are shown in Table 2. We calculated covariate shifts between the augmentation distribution P aug with the training P tr or test distribution P te . 'Aug-Train' and 'Aug-Test' represent GCS ( P aug , P tr ) and GCS ( P aug , P te ) , respectively. From the results in Table 2, we make these observations.

Discrepancy Principle. The term 'Original' denotes the training distribution prior to augmentation. It's observed that substantial covariate shifts exist between the training and test distributions, ranging from 0.419 to 0.557. The DropEdge technique notably expands Aug-Train , with a range of 0.627 to 0.851, while concurrently increasing Aug-Test , as evidenced by CMNIST (0.490 to 0.539) and Molbbbp (0.419 to 0.737). However, a distribution that deviates excessively from the test distribution may not effectively address the issue of covariate shift. FLAG, which perturbs only the node features, yields minor values in both Aug-Train and Aug-Test . G -Mixup notably augments AugTrain by generating OOD samples, but doesn't necessarily limit Aug-Test . Finally, AIA extends the disparity with the training distribution by augmenting environmental features, signifying that AIA can adeptly implement the principle of environmental feature discrepancy. Simultaneously, the imposed consistency constraint on stable features restricts the generated distribution from straying too far from the test distribution, as observed in Motif-base (0.557 to 0.462), Motif-size (0.522 to 0.098), and CMNIST (0.490 to 0.307).

Augmentation Diversity. We further delve into the diversity of data augmentation. The concept of augmentation diversity stems from the intuition that augmentations with greater degrees of freedom yield better performance [42]. Accordingly, we propose conditional entropy to

Motif

Molbbbp

Molhiv

55

60

65

70

75

Peformance

ERM

AIA

RDIA

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Stable feature ratio

55

60

65

70

75

0.01 0.1

0.2

0.5

1.0

1.5

2.0

3.0

5.0

Adversarial penalty

55

60

65

70

75

Motif (base)

Molbbbp (scaffold)

Molhiv (size)

measure the diversity of generated data: H ( ˜ G ∣ G ) = -E G [∑ ˜ g p ( ˜ g ∣ G ) log ( p ( ˜ g ∣ G ))] , where ˜ G and G represent the generated graph and original graph, respectively. To substantiate our ability to manage perturbed environmental features while preserving stable features, we examine diversity at the feature level, i.e., stable and environmental features. We employ the Motif dataset for validation due to its inclusion of ground-truth labels for stable and environmental features. The results in Table 3 reveal that our approach can guarantee the diversity of environmental features while constraining the variation of stable features.

To substantiate two principles of AIA, we present a selection of augmented graphs in Figure 3. These graphs are randomly sampled during the training phase. In the original Motif, the stable feature, is represented by the green portion and its type determines the label, while the yellow portion signifies the base-graph, or the environmental feature. Figure 3 ( Right ) exhibits the augmented samples generated during training. Nodes depicted in darker colors and edges with broader lines indicate higher soft-mask values. These results lead us to several noteworthy observations.

Visualization Analyses. AIA primarily perturbs the environmental features, while leaving the stable components undisturbed. In the Motif dataset, the base-graph represents a ladder and the motifgraph signifies a house . Following augmentation, the nodes and edges of the ladder graph undergo perturbations. However, the house component remains consistently stable throughout the training process. It shows that AIA successfully adheres to the proposed two principles, thus providing empirical support for our claims. Furthermore, under covariate shift, we also depict the stable features identified by AIA in comparison to other baseline methods (refer to Appendix E.4). It further underscores the limitations of alternative methods and highlights the superior performance of AIA.

5.4 Ablation Study (RQ3)

Adversarial Augmentation & Stable Feature Learning. As illustrated in Figure 4 ( Left ), 'w/o Adv' and 'w/o Sta' denote AIA without adversarial augmentation and without stable feature learning, respectively. RDIA is a variant that replaces adversarial augmentation in AIA with random augmentation ( i.e., random masks). The performance degrades when either component is used independently, compared to their combined application in AIA. The removal of adversarial perturbations results in a loss of the invariance condition inherent in stable feature learning [17, 29], leading to suboptimal outcomes. Conversely, the sole use of adversarial augmentation disrupts the stable features, thereby diminishing the performance. RDIA surpasses ERM, yet falls short of AIA, indicating that although randomness can foster discrepancy, it is less effective than the adversarial strategy.

Sensitivity Analysis. We conduct experiments to explore the sensitivities of ratio λ s and penalty coefficient γ . The results are displayed in Figure 4 ( Middle ) and ( Right ). λ s with 0.3 ∼ 0.8 performs well on Motif and Molbbbp, while Molhiv is better in 0.1 ∼ 0.3. It indicates that λ s is a dataset-sensitive hyper-parameter that needs careful tuning. For γ , the appropriate values range from 0.1 ∼ 1.5.

6 Related Work

Graph Data Augmentation [22, 23, 43] enlarges the training distribution by perturbing features in graphs. Recent studies [44, 13] observe that it often outperforms other generalization efforts [14, 31]. DropEdge [25] randomly removes edges, while FLAG [24] augments node features with an adversarial strategy. M-Mixup [26] interpolates graphs in semantic space. However, studies [14, 45]

point out that stable features are the key to OOD generalization. These augmentation efforts are prone to perturb the stable features, which easily lose control of the augmented data distribution.

Invariant Graph Learning has been widely adopted by recent works as a paradigm for handling distribution shifts on graphs. The pioneering works [3, 17] leverage the causal invariance principle to model the invariant predictive patterns in data for the OOD generalization purpose. With the similar spirit, GREA [18] and CAL [5] aim to learn stable features by considering different environments. Some other works also utilize invariant learning to develop generalizable models and algorithms for improving the generalization w.r.t. molecular graphs [8] and recommender systems [9].

Out-of-Distribution Learning on Graphs has aroused wide research interest in the graph learning community. One line of research is centered around the goal of improving the OOD generalization capabilities of models when encountered with test data from new unseen distributions [46, 3, 47, 39, 40, 17, 38, 48, 5]. Another line of research, differently, aims to identify the OOD data in the testing set and improve the reliablity of models against OOD data for which the model should reject for prediction [49, 50, 51]. The latter task is called Out-of-Distribution Detection in the literature and serves as another under-explored area that has different technical aspect from the present work. Due to space constraints, we put more discussions of other related studies in Appendix G.

7 Conclusion

In this study, we address the pervasive yet largely unexplored issue of covariate shift in graph learning. We introduce a novel graph augmentation method, AIA, grounded in two principles: environmental feature discrepancy and stable feature consistency. The discrepancy principle enables the model to explore new environments, thereby facilitating better generalization to potentially unseen environments. Meanwhile, the consistency principle maintains the integrity of stable features. We conduct extensive comparisons with various baseline models and perform thorough analyses.

8 Limitations and Broader Impacts

This paper presents a graph augmentation method, AIA, designed to bolster the academic community's application of data augmentation methodologies. We do not foresee any immediate, direct, or adverse societal implications resulting from our study's findings. We also present additional discussions regarding AIA's limitations and potential future work in Appendix H.

Acknowledgments and Disclosure of Funding

This research is supported by the National Natural Science Foundation of China (9227010114, U19A2079, 62302321) and the University Synergy Innovation Program of Anhui Province (GXXT2022-040). This work is also sponsored by Ant Group through CCF-Ant Research Fund and CCFAFSG Research Fund.

References

Zheyan Shen, Jiashuo Liu, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards out-of-distribution generalization: A survey. arXiv preprint arXiv:2108.13624 , 2021.

Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. Good: A graph out-of-distribution benchmark. arXiv preprint arXiv:2206.08452 , 2022.

Qitian Wu, Hengrui Zhang, Junchi Yan, and David Wipf. Handling distribution shifts on graphs: An invariance perspective. In ICLR , 2022.

Haoyang Li, Xin Wang, Ziwei Zhang, and Wenwu Zhu. Out-of-distribution generalization on graphs: A survey. arXiv preprint arXiv:2202.07987 , 2022.

Yongduo Sui, Xiang Wang, Jiancan Wu, Min Lin, Xiangnan He, and Tat-Seng Chua. Causal attention for interpretable and generalizable graph classification. In KDD , 2022.

Chenxiao Yang, Qitian Wu, Jiahua Wang, and Junchi Yan. Graph neural networks are inherently good generalizers: Insights by bridging GNNs and MLPs. In ICLR , 2023.

Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. G-mixup: Graph data augmentation for graph classification. In ICML , pages 8230-8248, 2022.

Nianzu Yang, Kaipeng Zeng, Qitian Wu, Xiaosong Jia, and Junchi Yan. Learning substructure invariance for out-of-distribution molecular representations. In NeurIPS , 2022.

Chenxiao Yang, Qitian Wu, Qingsong Wen, Zhiqiang Zhou, Liang Sun, and Junchi Yan. Towards out-of-distribution sequential event prediction: A causal treatment. In NeurIPS , 2022.

An Zhang, Jingnan Zheng, Xiang Wang, Yancheng Yuan, and Tat-Seng Chua. Invariant collaborative filtering to popularity distribution shift. In WWW , pages 1240-1251, 2023.

Yuan Gao, Xiang Wang, Xiangnan He, Zhenguang Liu, Huamin Feng, and Yongdong Zhang. Alleviating structural distribution shift in graph anomaly detection. In WSDM , pages 357-365, 2023.

Nanyang Ye, Kaican Li, Haoyue Bai, Runpeng Yu, Lanqing Hong, Fengwei Zhou, Zhenguo Li, and Jun Zhu. Ood-bench: Quantifying and understanding two dimensions of out-of-distribution generalization. In CVPR , pages 7947-7958, 2022.

Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre Alvise-Rebuffi, Ira Ktena, Taylan Cemgil, et al. A fine-grained analysis on distribution shift. In ICLR , 2022.

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019.

Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In ICML , pages 1448-1458, 2020.

Xingxuan Zhang, Peng Cui, Renzhe Xu, Linjun Zhou, Yue He, and Zheyan Shen. Deep stable learning for out-of-distribution generalization. In CVPR , pages 5372-5382, 2021.

Yingxin Wu, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. Discovering invariant rationales for graph neural networks. In ICLR , 2022.

Gang Liu, Tong Zhao, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Graph rationalization with environment-based augmentations. In KDD , pages 1069-1078, 2022.

Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science , 9(2):513-530, 2018.

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. In NeurIPS , 2020.

Shaohua Fan, Xiao Wang, Yanhu Mo, Chuan Shi, and Jian Tang. Debiasing graph neural networks via learning disentangled causal substructure. In NeurIPS , 2022.

Kaize Ding, Zhe Xu, Hanghang Tong, and Huan Liu. Data augmentation for deep graph learning: A survey. arXiv preprint arXiv:2202.08235 , 2022.

Tong Zhao, Gang Liu, Stephan Günnemann, and Meng Jiang. Graph data augmentation for graph machine learning: A survey. arXiv preprint arXiv:2202.08871 , 2022.

Kezhi Kong, Guohao Li, Mucong Ding, Zuxuan Wu, Chen Zhu, Bernard Ghanem, Gavin Taylor, and Tom Goldstein. Robust optimization as data augmentation for large-scale graphs. In CVPR , pages 60-69, 2022.

Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolutional networks on node classification. In ICLR , 2020.

Yiwei Wang, Wei Wang, Yuxuan Liang, Yujun Cai, and Bryan Hooi. Mixup for node and graph classification. In WWW , pages 3663-3674, 2021.

Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. In NeurIPS , 2020.

Susheel Suresh, Pan Li, Cong Hao, and Jennifer Neville. Adversarial graph augmentation to improve graph contrastive learning. In NeurIPS , pages 15920-15933, 2021.

Haoyang Li, Ziwei Zhang, Xin Wang, and Wenwu Zhu. Learning invariant graph representations for out-of-distribution generalization. In NeurIPS , 2022.

Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review. arXiv preprint arXiv:1908.05659 , 2019.

Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In ICLR , 2020.

Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. In ICLR , 2018.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In ICML , pages 214-223, 2017.

Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. Generalizing to unseen domains via adversarial data augmentation. NeurIPS , 31, 2018.

Alexey Dosovitskiy and Thomas Brox. Generating images with perceptual similarity metrics based on deep networks. NIPS , 2016.

Jaeho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. NeurIPS , 31, 2018.

David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In ICML , pages 5815-5826, 2021.

Siqi Miao, Miaoyuan Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic attention mechanism. In ICML , pages 15524-15543, 2022.

Haoyang Li, Xin Wang, Ziwei Zhang, and Wenwu Zhu. Ood-gnn: Out-of-distribution generalized graph neural network. IEEE TKDE , 2022.

Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. Generalizing graph neural networks on out-of-distribution graphs. arXiv preprint arXiv:2111.10657 , 2021.

Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, Kaili Ma, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng. Learning causally invariant representations for out-ofdistribution generalization on graphs. In NeurIPS , 2022.

Raphael Gontijo-Lopes, Sylvia Smullin, Ekin Dogus Cubuk, and Ethan Dyer. Tradeoffs in data augmentation: An empirical study. In ICLR , 2021.

Jaemin Yoo, Sooyeon Shim, and U Kang. Model-agnostic augmentation for accurate graph classification. In WWW , pages 1281-1291, 2022.

Mucong Ding, Kezhi Kong, Jiuhai Chen, John Kirchenbauer, Micah Goldblum, David Wipf, Furong Huang, and Tom Goldstein. A closer look at distribution shifts and out-of-distribution generalization on graphs. In NeurIPSW , 2021.

Chaochao Lu, Yuhuai Wu, José Miguel Hernández-Lobato, and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. In ICLR , 2021.

Qi Zhu, Natalia Ponomareva, Jiawei Han, and Bryan Perozzi. Shift-robust gnns: Overcoming the limitations of localized graph training data. NeurIPS , 34:27965-27977, 2021.

Wei Jin, Tong Zhao, Jiayuan Ding, Yozen Liu, Jiliang Tang, and Neil Shah. Empowering graph representation learning with test-time graph transformation. In ICLR , 2023.

Junchi Yu, Jian Liang, and Ran He. Mind the label shift of augmentation-based graph ood generalization. In CVPR , 2023.

Zenan Li, Qitian Wu, Fan Nie, and Junchi Yan. Graphde: A generative framework for debiased learning and out-of-distribution detection on graphs. In NeurIPS , 2022.

Qitian Wu, Yiting Chen, Chenxiao Yang, and Junchi Yan. Energy-based out-of-distribution detection for graph neural networks. In ICLR , 2023.

Yuxin Guo, Cheng Yang, Yuluo Chen, Jixi Liu, Chuan Shi, and Junping Du. A data-centric framework to endow graph neural networks with out-of-distribution detection ability. In KDD , 2023.

Masashi Sugiyama, Matthias Krauledat, and Klaus-Robert Müller. Covariate shift adaptation by importance weighted cross validation. Journal of Machine Learning Research , 8(5), 2007.

Zheyan Shen, Peng Cui, Kun Kuang, Bo Li, and Peixuan Chen. Causally regularized learning with agnostic data selection bias. In ACM MM , pages 411-419, 2018.

Renzhe Xu, Xingxuan Zhang, Zheyan Shen, Tong Zhang, and Peng Cui. A theoretical analysis on independence-driven importance weighting for covariate-shift generalization. In ICML , pages 24803-24829. PMLR, 2022.

Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. NeurIPS , 33:11996-12007, 2020.

Yue He, Xinwei Shen, Renzhe Xu, Tong Zhang, Yong Jiang, Wenchao Zou, and Peng Cui. Covariate-shift generalization via random sample weighting. AAAI , 2023.

Emanuel Parzen. On estimation of a probability density function and mode. The annals of mathematical statistics , 33(3):1065-1076, 1962.

Kurt Binder, Dieter Heermann, Lyle Roelofs, A John Mallinckrodt, and Susan McKay. Monte carlo simulation in statistical physics. Computers in Physics , 7(2):156-157, 1993.

Federico Monti, Davide Boscaini, Jonathan Masci, Emanuele Rodola, Jan Svoboda, and Michael M Bronstein. Geometric deep learning on graphs and manifolds using mixture model cnns. In CVPR , pages 5115-5124, 2017.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In ICLR , 2019.

Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR , 2017.

Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. In ICML , pages 1725-1735. PMLR, 2020.

Kiran K Thekumparampil, Chong Wang, Sewoong Oh, and Li-Jia Li. Attention-based graph neural network for semi-supervised learning. arXiv preprint arXiv:1803.03735 , 2018.

Christopher Morris, Nils M Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. In ICMLW , 2020.

Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In ICML , pages 6804-6814. PMLR, 2021.

Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful in an out-of-distribution generalization problem? arXiv preprint arXiv:2008.01883 , 2020.

Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In ICML , pages 145-155. PMLR, 2020.

Haohan Wang, Zeyi Huang, Xindi Wu, and Eric Xing. Toward learning robust and invariant representations with alignment regularization and data augmentation. In KDD , pages 1846-1856, 2022.

Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics , pages 4069-4077. PMLR, 2021.

Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research , 19(1):1309-1342, 2018.

Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2(11):665-673, 2020.

Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and Philip Yu. Generalizing to unseen domains: A survey on domain generalization. IEEE TKDE , 2022.

Jinheon Baek, Dong Bok Lee, and Sung Ju Hwang. Learning to extrapolate knowledge: Transductive few-shot out-of-graph link prediction. NeurIPS , 33:546-560, 2020.

Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. In ICML , pages 7313-7324. PMLR, 2021.

Massimiliano Mancini, Zeynep Akata, Elisa Ricci, and Barbara Caputo. Towards recognizing unseen categories in unseen domains. In ECCV , pages 466-483. Springer, 2020.

Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In ICML , pages 10-18. PMLR, 2013.

Yuning You, Tianlong Chen, Zhangyang Wang, and Yang Shen. Graph domain adaptation via theory-grounded spectral regularization. In ICLR , 2023.

Davide Buffelli, Pietro Lio, and Fabio Vandin. Sizeshiftreg: a regularization method for improving size-generalization in graph neural networks. In NeurIPS , 2022.

Gilad Yehudai, Ethan Fetaya, Eli Meirom, Gal Chechik, and Haggai Maron. From local structures to size generalization in graph neural networks. In ICML , pages 11975-11986. PMLR, 2021.

Zeyang Zhang, Xin Wang, Ziwei Zhang, Haoyang Li, Zhou Qin, and Wenwu Zhu. Dynamic graph neural networks under spatio-temporal distribution shift. In NeurIPS , 2022.

Junfeng Fang, Xiang Wang, An Zhang, Zemin Liu, Xiangnan He, and Tat-Seng Chua. Cooperative explanations of graph neural networks. In WSDM , pages 616-624, 2023.

Junfeng Fang, Wei Liu, An Zhang, Xiang Wang, Xiangnan He, Kun Wang, and Tat-Seng Chua. On regularization for explaining graph neural networks: An information theory perspective. 2022.

Yongduo Sui, Xiang Wang, Tianlong Chen, Meng Wang, Xiangnan He, and Tat-Seng Chua. Inductive lottery ticket learning for graph neural networks. Journal of Computer Science and Technology , 2023.

Tianlong Chen, Yongduo Sui, Xuxi Chen, Aston Zhang, and Zhangyang Wang. A unified lottery ticket hypothesis for graph neural networks. In ICML , pages 1695-1706. PMLR, 2021.

Yanfang Wang, Yongduo Sui, Xiang Wang, Zhenguang Liu, and Xiangnan He. Exploring lottery ticket hypothesis in media recommender systems. International Journal of Intelligent Systems , 37(5):3006-3024, 2022.

Yuan Gao, Xiang Wang, Xiangnan He, Zhenguang Liu, Huamin Feng, and Yongdong Zhang. Addressing heterophily in graph anomaly detection: A perspective of graph spectrum. In WWW , pages 1528-1538, 2023.

Yuan Gao, Xiang Wang, Xiangnan He, Huamin Feng, and Yong-Dong Zhang. Rumor detection with self-supervised learning on texts and social graph. Frontiers Comput. Sci. , 17(4):174611, 2023.

Yanqiao Zhu, Yuanqi Du, Yinkai Wang, Yichen Xu, Jieyu Zhang, Qiang Liu, and Shu Wu. A survey on deep graph generation: Methods and applications. arXiv preprint arXiv:2203.06714 , 2022.

A Correlation Shift and Covariate Shift

Following graph data generation process [17, 8, 3, 29], we can observe that environmental features easily change outside the training distribution, owing to their unstable nature. Hence, distribution shifts are mainly caused by the environmental features [12, 13, 2]. Specifically, we define the joint distribution of training and test data as P tr ( G,Y ) and P te ( G,Y ) , respectively. Since their joint distribution can be rewritten as P tr ( G,Y ) = P tr ( Y ∣ G ) P tr ( G ) and P te ( G,Y ) = P te ( Y ∣ G ) P te ( G ) , we can find that there exist two main reasons that lead to the distribution shift P tr ( G,Y ) ≠ P te ( G,Y ) . We give intuitive examples in Figure 5 and formal definitions of these two distribution shifts.

Correlation shift P tr ( Y ∣ G ) ≠ P te ( Y ∣ G ) , P tr ( G ) = P te ( G ) . If the statistical correlation of environmental features and labels is inconsistent in training and test data, a well-fitted model in training data may fail in test data, which is also known as spurious correlation [14], correlation shift [12] or concept shift [2]. Formally, correlation shift describes the conditional distribution P tr ( Y ∣ G ) ≠ P te ( Y ∣ G ) .

Covariate shift P tr ( G ) ≠ P te ( G ) , P tr ( Y ∣ G ) = P te ( Y ∣ G ) . If there exist environmental features in the test distribution that the model has not seen during training, it will also result in a large performance drop. This unseen distribution shift is well known as covariate shift [2] or diversity shift [12]. It means that the environmental features in test data are unseen in training data, which leads to P tr ( G ) ≠ P te ( G ) . In Definition 2.1, we also quantitatively measure the covariate shift between P tr ( G ) and P te ( G ) .

It is worth noting that within the computer vision domain, the general covariate shift is frequently synonymous with sample selection bias [52, 53, 54, 55, 56]. Various factors contribute to covariate shift, such as heterogeneous category distribution or domain-specific variances. In the context of our investigation into graph-based models, we adhere to the assumptions outlined in previous literature [2, 12], which mainly ignore the influence of label shifts. And we posit that covariate shifts are principally attributed to the environmental features. The exploration of more comprehensive scenarios involving covariate shifts will be undertaken in our future work.

…

Training data

Test data

Correlation Shift:

𝑃𝑡𝑟 (𝑌|𝐺) ≠ 𝑃 𝑡𝑒(𝑌|𝐺)

𝑒1

𝑒2

𝑒1

𝑒2

𝑒1

𝑒2

𝑒1

𝑒2

Training data

Test data

Covariate Shift:

𝑃𝑡𝑟 (𝐺) ≠ 𝑃 𝑡𝑒(𝐺)

𝑒1

𝑒2

𝑒1

𝑒2

𝑒1

𝑒2

𝑒1

𝑒2

Environmental feature

Stable feature

𝑌 1

:

House

𝑌 2

:

Cycle

𝑒1 :

Tree

𝑒2 :

Ladder

Data Generation Mechanism

…

…

…

…

…

…

…

B Proofs

In this section, we provide the detailed proofs to our proposition. We start with the definition of the 1-Wasserstein distance, D W ( P,P ′ ) , between two distributions P and P ′ :

where Γ ( P,P ′ ) represents the set of all joint distributions π ( x , y ) that have P and P ′ as their respective marginals. Under our conditions, due to P ′ ( x -x ′ ) = P ( x ) , we can precisely determine the manner in which the mass from P was transferred to create P ′ . This process involves moving each point x under P to x + x ′ under P ′ . Therefore, the infimum is attained by the coupling that deterministically transitions each point x to x + x ′ . Consequently, the Wasserstein distance simplifies to:

This result establishes that the Wasserstein distance between the distributions P and P ′ is equivalent to the magnitude of x ′ . Therefore, we obtain:

Similarly, D W ( P, ̃ P 2 ) = ∣∣ x 2 ∣∣ . Consequently, if D W ( P, ̃ P 1 ) < D W ( P, ̃ P 2 ) , we can deduce that ∣∣ x 1 ∣∣ < ∣∣ x 2 ∣∣ .

Next, we examine the GCS measure. The covariate shift between P and ̃ P 1 is given by:

where S 1 and S 2 denote the non-overlapping regions between P and ̃ P 1 , and between P and ̃ P 2 , respectively. Now we define another ̃ P ′ 1 that satisfies ̃ P ′ 1 ( x ) = P ( x +∣∣ x 1 ∣∣ ⋅ x 2 ∣∣ x 2 ∣∣ ) and let S ′ 1 denote the non-overlapping regions between P and ̃ P ′ 1 . Given the isotropy of ̃ P 1 , the integral over S 1 is equal to the integral over S ′ 1 . Therefore, we can proceed to deduce:

Let S ′ 2 = { x +(∣∣ x 2 ∣∣ - ∣∣ x 1 ∣∣) ⋅ x 2 ∣∣ x 2 ∣∣ ∣ x ∈ S ′ 1 } , we can easily know that ∣S ′ 1 ∣ = ∣S ′ 2 ∣ and S ′ 2 ⊆ S 2 . Thus we have:

The integral over S 2 could not be less than the integral over S 1 , leading to GCS ( P, ̃ P 1 ) ≤ GCS ( P, ̃ P 2 ) . Hence, we complete the proof.

C Estimation of Graph Covariate Shift

In this section, we elaborate on the implementation details of estimating the graph covariate shift. Without loss of generality, we start with the example of estimating the graph covariate shift between the training and test distributions. Given the training set and test set D tr and D te , they follow probability distribution functions P tr and P te . The process of estimating GCS ( P tr , P te ) is summarized in the following two steps:

Firstly, it is intractable to directly estimate the distribution in graph space G . Inspired by [12], we can obtain the graph features and estimate the distribution in feature space F . Specifically, given a sample, we train a binary GNN classifier f to distinguish which distribution it comes from, where f (⋅) = Φ ○ h , h (⋅) ∶ G → F is a graph encoder, and Φ (⋅) ∶ F →{ 0 , 1 } is a binary classifier. Then we can adopt the pre-trained GNN encoder h to extract graph features.

Secondly, we prepare the features and estimate the distribution of the data via Kernel Density Estimation (KDE) [57]. Finally, we adopt the Monte Carlo Integration under importance sampling [58] to approximate the integrals in Definition 2.1.

Similarly, we have:

We summarize these implementations in Algorithm 1. In lines 4 and 5, to avoid the label shift [12], we adopt sample reweighting to ensure the balance of each class.

Algorithm 1: Estimation of Graph Covariate shift

Require: Training dataset D tr and test dataset D te ; Batch size N ; Loss function ℓ ; GNN f = Φ ○ h ; Importance sampling size M ; Threshold ϵ . Ensure: Estimated covariate shift GCS ( P tr , P te ) . 1: Initialize parameters of f 2: # Train a graph classifier 3: while not converge do 4: Sample a batch B tr ←{( g i , y i )} N i = 1 ⊂ D tr and relabel all y i ← 0 5: Sample a batch B te ←{( g i , y i )} N i = 1 ⊂ D te and relabel all y i ← 1 6: B ← B tr ∪B te 7: for each ( g i , y i ) ∈ B do 8: Compute loss ℓ ( f ( g i ) , y i ) and back-propagate gradients 9: end for 10: Update the parameters of f via gradient descent and reset the gradients 11: end while 12: # Prepare the features for the estimation 13: Extract training and test feature sets F tr and F te via encoder h 14: F ← F tr ∪F te 15: Scale F to zero mean and unit variance 16: ˆ ω ← fit by KDE the distribution of F 17: Split F to recover the original partition F ′ tr , F ′ te 18: ˆ P tr , ˆ P te ← fit by KDE the distributions of F ′ tr , F ′ te 19: # Estimate the covariate shift 20: Initialize GCS ( P tr , P te ) ← 0 21: for t ←{ 1 , ..., M } do 22: z ← sample from ˆ ω 23: if ˆ P tr ( z ) < ϵ or ˆ P te ( z ) < ϵ then 24: GCS ( P tr , P te ) ← GCS ( P tr , P te ) + ∣ ˆ P tr ( z ) -ˆ P te ( z )∣/ ˆ ω ( z ) 25: end if 26: end for 27: GCS ( P tr , P te ) ← GCS ( P tr , P te )/ 2 M

Algorithm 2: Adversarial Invariant Augmentation

Require: Training set D tr ; Adversarial augmenter T θ 1 (⋅) ; Stable feature generator T θ 2 (⋅) ; GNN classifier f (⋅) with parameters θ ; Learning rates α,β ; Batch size N ; Stable feature ratio λ s ; Penalty γ . 1: Randomly initilize θ , θ 1 , θ 2 2: while not converge do 3: Sample a batch B tr ←{( g i , y i )} N i = 1 ⊂ D tr 4: for each ( g i , y i ) ∈ B tr do 5: M a adv , M x adv ← T θ 1 ( g i ) // adversarial perturbations 6: M a sta , M x sta ← T θ 2 ( g i ) // regions of stable features 7: ̃ M a ←( 1 -M a sta ) ⊙ M a adv + M a sta // augment edges 8: ̃ M x ←( 1 -M x sta ) ⊙ M x adv + M x sta // augment nodes 9: ̃ g i ←( A i ⊙ ̃ M a , X i ⊙ ̃ M x ) // augmented graph 10: end for 11: Compute L adv -L reg 1 via Equation (7) and (9) 12: Compute L sta +L reg 2 via Equation (8) and (10) 13: Update parameters of adversarial augmenter via gradient ascent: θ 1 ← θ 1 + α ∇ θ 1 (L adv -L reg 1 ) 14: Update parameters of GNN and stable feature generator via gradient descent: θ ← θ -β ∇ θ (L sta +L reg 2 ) ; θ 2 ← θ 2 -β ∇ θ 2 (L sta +L reg 2 )

15: end while

D Implementation Details

D.1 Algorithm

We summarize the detailed implementations of AIA in Algorithm 2. We alternately optimize the adversarial augmenter and stable feature generator with the backbone model, in lines 13 and 14. We adopt the learned stable features for predictions in the inference stage.

D.2 Datasets

In this paper, we conduct experiments on graph OOD datasets [2] and OGB datasets [20], which include Motif, CMNIST, Molbbbp and Molhiv. We follow [2] to create various covariate shifts, according to base, color, size and scaffold splitting. Base, color, size and scaffold are features of the graph data and do not determine the labels of the data, so they can be regarded as environmental features. The statistics of the datasets are summarized in Table 4. Below we give a brief introduction to each dataset.

Motif: It is a synthetic dataset from Spurious-Motif [17, 5]. As shown in original graphs in Figure 5, each graph is composed of a base-graph ( wheel, tree, ladder, star, path ) and a motif ( house, cycle, crane ). The label is only determined by the type of motif. We create covariate shift according to the base-graph type and the graph size ( i.e., node number). For base covariate shift, we adopt graphs with wheel, tree, ladder base-graphs for training, star for validation and path for testing. For size covariate shift, we use small-size of graphs for training, while the validation and the test sets include the middle- and the large-size graphs, respectively.

CMNIST: Color MNIST dataset contains graphs transformed from MNIST via superpixel techniques [59]. We define color as the environmental features to create the covariate shift. Specifically, we color digits with 7 different colors, where five of them are adopted for training while the remaining two are used for validation and testing.

Molbbbp & Molhiv: These are molecular datasets collected from MoleculeNet [19]. We define the scaffold and graph size ( i.e., node number) as the environmental features to create two types of covariate shifts. For scaffold shift, we follow [2] and use scaffold split to create training, validation

and test sets. For size shift, we adopt the large-size of graphs for training and the smaller ones for validation and testing.

D.3 Metrics

We adopt classification accuracy as the metric for Motif and CMNIST. As suggested by [20], we use ROC-AUC for Molhiv and Molbbbp datasets. In addition, we use GCS ( P,Q ) to measure the covariate shift between distributions P and Q . For all experimental results, we perform 10 random runs and report the mean and standard derivations. For augmentation diversity, we use conditional entropy to measure the diversity of generated data. We normalize it to [ 0 , 1 ] for better comparison. Specifically, for a given graph data, we compute the conditional entropy by collecting augmented data generated during the training process. We conduct experiments by collecting 1000 graph data, and report the mean and standard deviation.

D.4 Training Settings

We use the NVIDIA GeForce RTX 3090 (24GB GPU) to conduct all our experiments. To make a fair comparison, we adopt GIN [60] as the default architecture to conduct all experiments. We tune the hyper-parameters in the following ranges: α and β ∈ { 0 . 01 , 0 . 005 , 0 . 001 } ; λ 2 ∈ { 0 . 1 , ..., 0 . 9 } ; γ ∈ { 0 . 01 , 0 . 1 , 0 . 2 , 0 . 5 , 1 . 0 , 1 . 5 , 2 . 0 , 3 . 0 , 5 . 0 } . The hyper-parameters are summarized in Table 5.

D.5 Baseline Settings

For a more comprehensive comparison, we selected 16 baselines. In this section, we give a detailed introduction to the settings of these methods.

For ERM, IRM [14], GroupDRO [31], VREx [37], and M-Mixup [26], we report the results from the study [2] by default and reproduce the missing results on Molbbbp.

For DIR [17], CAL [5], GSAT [38], DropEdge [25], GREA [18], FLAG [24], G -Mixup [7], CIGA [41] and DisC [21], they provide source codes for the implementations. We adopt default settings from their source codes and detailed hyper-parameters from their original papers.

For OOD-GNN [39] and StableGNN [40], their source codes are not publicly available. We reproduce them based on the codes of StableNet [16].

For RDIA in Section 5.4, it is a variant that replaces the adversarial augmentation in AIA with random augmentation. In our implementation, we use all-one matrices to create the initial node and edges masks. Then we randomly set 20% of nonzero elements to zero in these masks. Finally, we apply these masks to the graphs for random data augmentation. The process of stable feature learning is consistent with AIA.

E More Experimental Results

E.1 Results on Correlation Shift

Original graphs

AIA (Ours)

GREA

DropEdge

CAL

Although this work focuses on the OOD issue of covariate shift, for completeness, we also evaluate the performance of AIA under correlation shift. Following [2], we choose three graph OOD datasets ( i.e., Motif, CMNIST, Molhiv) with three different graph features ( i.e., base, color, size) to create correlation shifts. For baselines, we choose three generalization algorithms ( i.e., ERM, IRM [14], VREx [37]), three graph generalization methods ( i.e., DIR [17], CAL [5], OOD-GNN [39]) and three data augmentation methods ( i.e., DropEdge [25], FLAG [24], M-

Mixup [26]). The experimental results are shown in Table 7. We can observe that AIA can also effectively alleviate the correlation shift. These results demonstrate that AIA learns better stable features by encouraging environmental discrepancy, which can effectively break spurious correlations that are hidden in the training data.

E.2 Results on Diverse Backbones

We select three different GNN backbone models (GCN [61], GCNII [62] and GAT [63]) for experiments. From the results in Table 6, our observations and conclusions remain the same with the diverse backbones.

E.3 Results on More Real-world Datasets

To demonstrate the effectiveness of the proposed AIA, we also conduct experiments on commonly used TU datasets [64], which include MUTAG, NCI1, PROTEINS, COLLAB, IMDB-B, IMDBM. For training settings, we follow CAL [5] and adopt GIN [60] as our backbone model. The experimental results are shown in Table 8. For the results, we can observe that our method can achieve the best performance over different datasets.

E.4 More Visualizations

To demonstrate the superiority of our method, we also visualize the captured stable features by AIA and compare them with other baselines. The results are displayed in Figure 6. From the results, we can easily observe that our method can find stable parts more accurately than other baseline methods.

F Complexity Analyses

Firstly, we define the average numbers of nodes and edges per graph in the dataset to be n and m , respectively. Let N denote the batch size, l , l a and l c denote the numbers of layers in the GNN backbone, adversarial augmenter and stable feature generator, respectively. d , d a and d c are the

dimensions of hidden layers in the GNN backbone, adversarial augmenter and stable feature generator, respectively.

Time Complexity. The time complexity of the adversarial learning objective is O( N ( l a md a + 2 lmd )) . For the stable feature learning objective, the time complexity is O( N ( l c md c + 2 lmd )) . For the regularization terms, the time complexity is O( 2 N ( n + m )) . For simplicity, we assume l a = l c and d a = d c . Hence, the time complexity of a forward propagation is O( 2 N ( l a md a + 2 lmd + n + m )) .

Model Size. In addition to the GNN backbone model, we also introduce two small networks for adversarial augmentation and stable feature learning. In our implementations, the parameters of AIA are around twice as large as those of the original GNN model. The running time and model size are shown in Table 9.

G More Related Studies

OODGeneralization [1, 65, 66] has been widely explored. Inspired by invariant learning and causal theory, a series of general algorithms [67, 68, 69, 66, 70] have recently been proposed to solve the OOD problem. Recent studies [12, 2, 13] point out that OOD falls into two specific categories: correlation shift ( aka. concept shift) and covariate shift ( aka. diversity shift). Correlation shift denotes that the environmental features and labels establish a statistical correlation that is inconsistent in training and test data. Thus, the models prefer to learn spurious correlations and rely on shortcut features [71] for predictions, resulting in a large performance drop. In contrast, covariate shift indicates that there exist unseen environmental features in test data. The limited training environment makes this issue intractable. Tasks pertaining to domain generalization [72] often exhibit the covariate shift issue, such as model inference under previously unseen test domains. Consequently, there have been numerous recent efforts [73, 74, 75, 76] to address such challenges. In recent years, OOD generalization on graphs is drawing widespread attention [4]. Given the intricate nature of graph data types and associated tasks [77], issues related to distribution shifts can emerge in various contexts, such as node classification [46, 3, 47], graph classification [39, 40, 17, 38, 48, 5, 18, 78, 79], or dynamic graph data [80], among others. Emerging research has increasingly focused on the identification of stable features or subgraphs within graph data. These substructures are posited to maintain causal relationships with target labels, thereby enhancing the model's capacity for generalization [5, 18], explainability [81, 17, 38, 82], and computational efficiency [83, 84, 85]. Concurrently, efforts are being made to extend these principles to diverse applications, such as molecular graphs [8], recommender systems [9, 10], and anomaly detection [11, 86, 87].

Comprehensive Comparisons with EERM [3]. Although EERM shares similar goals with us, generating several environments through augmentation, there exist many technical and contribution differences. Firstly, EERM ignores the distinction between correlation shift and covariate shift problems, while we distinguish these two shifts in detail and design a framework specifically for covariate shift. Secondly, EERM does not model stable and environmental features, which results in the inability to explicitly distinguish them. In contrast, we explicitly model the environmental and stable features. Hence, we can effectively identify stable and environmental features and explicitly

separate them from data. Thirdly, we also design a metric, GCS ( ̃ P,P ) , which can effectively measure the discrepancy of the environmental features for our augmented data. And we directly encourage the environmental discrepancy of the augmented samples by maximizing GCS ( ̃ P,P ) . However, EERM does not provide any evaluation metric for environmental discrepancy. To encourage the discrepancy, they 'blindly' maximize the variance of the empirical risk in K environments. Finally, for generalization scope, EERM is based on the IRM [14] by minimizing the empirical risk in K environments. In contrast, inspired by DRO [31], we can guarantee the generalization within the robust radius ρ . We summarize the above detailed discussions in Table 10.

H Limitation & Future Work

Although AIA outperforms numerous baselines and can achieve outstanding performance under various covariate shifts, we also prudently introspect the following limitations of our method. And we leave the improvements of these limitations as our future work.

AIA performs OOD exploration through an adversarial data augmentation strategy to achieve environmental discrepancy. However, it only perturbs the existing graph data in a given training set, such as perturbing original graph node features or graph structures. Hence, it is possible that there still exist some overlaps between the augmented distribution and training distribution, so discrepancy principle cannot be thoroughly achieved. In future work, we will attempt to design more advanced data augmentation methods, such as graph generation-based strategies [88], to generate more unseen and novel graph data, for pursuing the discrepancy principle.

For model training, we adopt adversarial training and stable feature learning to alternately optimize the adversarial augmenter, stable feature generator and backbone GNN. This training strategy may make the training process unstable, so the performance of AIA may experience a large variance. In addition, these two networks also involve additional parameters. Optimizing these parameters separately will also increase the time complexity, as shown in Appendix F. Hence, in future work, we will explore how to utilize more advanced optimization methods and lightweight models to achieve the principles of environmental feature discrepancy and stable feature consistency.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an augmentation technique to handle covariate shift during graph classification by creating new \"environmental factors\" and simultaneously preserving \"stable\" factors. Since unseen environmental factors occur during covariate shift, the authors propose using an adversarial augmentor to find augmentations that increase the \"GCS.\" However, since unconstrained adversarial augmentation could harm the \"stable\" factors, the method includes a second network that learns to a mask over the stable nodes/edges. The found mask is used to help keep the adversarial augmentation to environmental factors and also in the loss. The method is shown to do well on the a variety of benchmarks. \n\n**POST REBUTTAL:** I have read the other reviews and the author's rebuttals. I thank the authors for their thorough responses! The explanations and additional experiments (non-random augs., graph transformers) are helpful and strengthen the paper. I encourage the authors to add these discussions, and experiments to their final paper as well as cite the suggested related work on task-relevant invariances. My concerns have been addressed and have raised the score!\n\nSTRENGTHS:\n- Motivation: The method is well-motivated: the environmental feature discrepancy and stable feature consistency objectives for promoting performance under covariate shift make sense. \n\n- Strong experimental results. The proposed method is evaluated against many baselines on several different benchmarks. The evaluation of augmentation diversity and discrepancy is also useful.\n\nWEAKNESSES:\n- Limited Novelty: Individual pieces, as well as the motivation for the method are not particularly novel. These concepts have been discussed in the context of designing augmentations in general and also in the context of creating augmentations for SSL (in-variance/recover-ability, style/content). Adversarial augmentation has been explored by other methods, and masking for explanation too. To this end, I don't think that the theoretical discussion adds too much. \n\n- Complexity: Including two networks, as well as an adversarial objective is expensive and potentially more difficult to train. Moreover, the masking operation can be expensive, since the mask must be learned over all the nodes and the edges.\n\nQUESTIONS:\n- The discussion of augmentation diversity in Table 3 is interesting. However, I was wondering if it would be possible to also compare to one of the augmentation methods too that is not random? Curious to see if how well something like FLAG compares. \n\n- Quick (not important) question: how where hyper-parameters selected during tuning? \n\n- Quick question: can you add some run-times/model sizes to the appendix for later drafts?\n\nLIMITATIONS:\nSee above. Negative societal impacts, etc are adequately discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies on the covariate shift problem in graph classification tasks, as opposed to the more commonly studied correlation shift. The authors propose a novel data augmentation strategy, Adversarial Invariant Augmentation (AIA), guided by two principles, environmental feature discrepancy and stable feature consistency. The method leverages an adversarial augmenter to adversarial generate masks. Another stable feature generator is also utilized to promote stable feature consistency. The proposed approach equips the graph classifier model with an enhanced ability to identify stable features in new environments and effectively mitigate the covariate shift issue. \n\n**Post rebuttal**: I'm content with the rebuttal of the authors and will keep my rating.\n\nSTRENGTHS:\n- This paper addresses the under-explored perspective on graph OOD, which is covariate shift. It provides new insights into the important OOD area of graph learning.\n- This paper designs the AIA methodology corresponding with the two proposed principles. The adversarial augmentation method design is novel with good intuitions. \n- The experimental results are extensive covering a list of baseline methods and datasets with diverse properties. AIA outperforms specialized graph generalization and augmentation algorithms. \n- The authors provide in-depth analysis discussing experiments with respect to the principles, which is persuasive. \n- Ablation study demonstrates the effectiveness of the designed components.\n- The structure of the paper is clear. The writing is easy to follow.\n\nWEAKNESSES:\n- Advsersarial training is notoriously inefficient at neural network training, so leveraging adversarial training to solve the graph OOD problem may affect the training/convergence speed of GNNs. It would be better if authors can provide discussions regarding such overhead.\n\nQUESTIONS:\n- What is the GNN backbone used in AIA for the experiments?\n- How will the choice of GNN backbones affect the final outcome?\n\nLIMITATIONS:\nThe paper addresses the limitations in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a data augmentation technique on graph datasets. The model consists of two main components: adversarial augmenter and stable feature generator. The adversarial augmenter tries to generate new samples by generating dropping masks for the nodes and edges adversarially, within some augmentation cost. The stable feature generator, on the other hand, tries to identify subsets of graphs that are preserved among all the samples. The authors argue that this construction is more suitable for covariate shift settings compared to the previous graph augmentation methods like DropEdge. Finally, the authors demonstrate the effectiveness of the technique in the experiments.\n\nSTRENGTHS:\n- Data augmentations on graphs is an important area that needs more innovations, as many prominent techniques for data augmentation in other domains, like vision, do not apply to graph data. \n- The paper provides explanation details on covariate shift problem in graph setting\n- The technique proposed by the authors is interesting. The use of two components: the stable feature generation and adversarial augmenter, on graph data, is new. \n- The technique trains the GNN, stable feature generation, and adversarial augmenter together in an alternative way.\n- The results demonstrate that the proposed technique outperforms baselines on real-world covariate shift datasets.\n\nWEAKNESSES:\n- First, I would like to clarify the authors' claim that \"covariate shift is less explored type of OOD\" in the introduction. Contrary to this claim, covariate shift (sometimes also called sample selection bias) is a well-studied problem with a tremendous amount of previous and current research works. The efforts to develop algorithms to correct the shift/bias have been made by many previous works. The most notable family of techniques is the important weighting technique. I suggest the author deep dive more into the literature about covariate shift/sample selection bias.\n- There are some unclear descriptions of the methods that I would like the author to clarify. Particularly, it's not clear to me how the stable feature generator can achieve the desired results of finding features (or subgraphs) that persist in most graphs in the dataset. I understand its mask construction tries to shelter some of the components in the graph from being perturbed by the augmenter. However, from the loss function, it is not clear how this construction end-up generating mask for common patterns in the dataset, not just creating a unique mask for each graph sample. \n- The connection of the paper to covariate shift is also a bit misleading. In the standard covariate shift/sample selection bias setting, there are no need or stable features to exist in the dataset. The covariate shift only requires the sample distribution P(x) between training and testing to be different, while the label conditional distribution P(y|x) remains the same (Shimodaira, 2000; Sugiyama et.al, 2007). The testing sample distribution may or may not have some common characteristics with the training distribution.\n- The paper assumes the existence of stable features and actively finds them in the training dataset. In the standard covariate shift setting, even though some of the characteristics commonly occur in the training data, there is no guarantee that they also occur in the test data.\n- In short, the authors provide an interesting technique for data augmentation on graph data and demonstrate its effectiveness. However, attributing it to the covariate shift setting is misleading.\n\nRef: \n- H. Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of Statistical Planning and Inference, 90(2):227–244, 2000\n- Sugiyama, M., Krauledat, M., & Müller, K. R. (2007). Covariate shift adaptation by importance weighted cross validation. Journal of Machine Learning Research, 8(5).\n\nQUESTIONS:\nPlease answer my concerns in the previous section.\n\nLIMITATIONS:\nNo concern.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to enable GNNs handle covariate shifts through the use of synthetic augmentations. An inherent challenge in construction task-specific augmentations is to appropriately handle style (or environment features) and content (or stable features). While other existing works have also emphasized the need to promote task-relevant invariances (e.g., Analyzing data-centric properties for graph contrastive learning, Neurips 2022), this paper focuses on automatically identifying those features through an adversarial training strategy. Results on graph classification benchmarks demonstrate the benefits of AIA.\n\n**Post Author Rebuttal**:The authors have reasonably responded to all my concerns. Hence, I raise my score!\n\nSTRENGTHS:\n+ The paper is well written and easy to follow. The problem is clearly laid out and the I personally liked the organization of the experiments in the form of research questions.\n+ Though the proposed algorithm builds upon existing formalisms and ideas from the certified robustness literature, I believe the extension and its implementation to this problem is sufficiently novel.\n+ The theoretical analysis is intuitive and well presented.\n+ Experiment results and the hyperparameter study provide a convincing demonstration of the proposed approach.\n\nWEAKNESSES:\n1. At the outset, the idea of splitting the adversarial augmenter and stable feature generator with independent (learnable) masks appears challenging to solve. There is a risk that the same entries can be picked as relevant in both the masks. While the regularizer checks for the ratio, will it benefit to include an explicit constraint to make the masks disjoint? This is often done in state-of-the-art source separation algorithms that aim to split the different sources from a given observation.\n2. It will be beneficial if the authors can intuitively explain how this approach is able to handle size generalization.\n3. More insights into the results will help. For example on the Motif benchmark, when compared to the baselines, the bigger benefits are observed in the base setting. Why is that?\n4. While GIN has been used for all experiments, will the benefits persist with a graph transformer?\n\nQUESTIONS:\nsee weaknesses\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}