Unleashing the Power of Graph Data Augmentation on Covariate Distribution Shift
Authors: Yongduo Sui
Qitian Wu
Jiancan Wu
Qing Cui
Longfei Li
JUN ZHOU
Xiang Wang
Xiangnan He
Venue: NeurIPS 2023 poster (2023)
Keywords: Graph Neural Network
Graph Data Augmentation
Distribution Shift
TL;DR: We propose a novel graph data augmentation method to address covariate distribution shifts.

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
Yongduo Sui Qitian Wu?, Jiancan Wu', Qing Cui², Longfei Li², Jun Zhou", Xiang Wang", Xiangnan He University of Science and Technology of China, Shanghai Jiao Tong University, *Ant Group syd20190maii ustc. edu en, echo7400sjtu, edu {wujcan, xiangwang 1223, xiingnanne) {cuiqing eq, longyao. 11f jun. zun /houjun/Oantgrouup Santgroupt com
edu, cn, com,
Abstract
The issue of distribution shifts is emerging as critical concern in graph represen- tation 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 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 perturbed stable features, thus exposing limitations in handling the problem of covariate shift. To address this challenge, propose simple-yet-effectivy data augmentation strategy, Ad- versarial 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 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 implementa- tion codes are publicly available uthhtps::/ gighub.......... com http:ooooooooooooo yongduouiiiAAAA
Introduction
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 [I] 2]. However, this as- sumption often falls short in real-world scenarios due to the out-of-distribution (OOD) data that potentially exists during the testing phase [ which results in distribution shifts between training and test graphs. tion 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
result, there is increasing research interest OOD generaliza-
This work done during author's 's internship Ant Group. "Corresponding authors. Xiang Wang and Xiangnan He also affiliated with Institute of Artificial Intelligence, Institute of Dataspace, Hefei Comprehensive National Science Center.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

With more specific views, other attempts focus on designing generalizable models for tackling distribu- tion shifts in particular domains with certain data for- mats, egg,, molecular graphs [ recommender sys-
Pirsin
P Parop
Dropkdga
gasts
tems p 10 and anomaly detection II]
Covariate shift
Pouur
However, the majority of existing studies primarily fo- cus on the correlation shift, one type of distribution shift concerning OOD generalization [2 13,, leav- ing another equally important type of distribution shift, i..., the covariate shift, largely under-explored in graph representation learning. From the perspective of invari- ant 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 hy- there often exist stable features, which are informativ features of the entire graphs and can reflect the predictive patterns in data. Based 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 distribu- tion shifts: (1) correlation shift indicates that environments and labels establish inconsistent statistical correlations in training and test data, under the assumption that test environments covered in the training data; whereas, (2) covariate shift characterizes that the environmental features in test data are unseen in training data 12)..
Test
AlA popts
grapas
Prest
stable Sesture (toraee)
Envinnmertal festume (laditee tree)
Figure 1: Perain and Pest denote the train- ing and test distributions. Párop and Pours represent the distributions of augmented data via DropEdge and AIA.
pothesis in prior studies [17 [8 5
Considering a toy example in Figure T 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:)) are stable to determine the predictive property of molecules [5] 8). Whereas, scaffolds (e.g., carbon rings) are usually patterss irrelevant to the molecule properties, which can be seen as environmental features [2] [9]. In practice, 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
Considering the differences between correlation and covariate shifts, we take 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 Gruph 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. can be roughly divided into node-level [24] edge-level [25), and graph-level (26) 7 with random 277 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 T 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 fea- ture 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 D ATA 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 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.. b4)} denote a dataset that is divided into training set Dur { } and a test set Die {9g yf JesE... Eer and E. are the index sets of training and test environments, respectively.
2.1 Definitions and Problem Formations
In this work, focus on the graph classification scenario, which aims train models with Dir and infer the labels in Die. such as molecular property prediction. From the viewpoints of invariant learning and stable learning, the inner mechanism of the labeling rule G + is usually assumed to depend on the stable features [17 I8
5 g 21]. which are informative substructures of the entire graph. The relationship between the stable features and labels 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 [7 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 envirommental features and the limitations in the data collection process. the training and test distributions are often inconsistent in real-world scenarios, i.e., Par G, Y) # Pie (G, Y). which leads two main types of distribution shifts 2 12: (1) Correlation shift (aka. spurious correlation or concept shift [2) refers to (G)Y Po.(Gl)) Pir(G) Pte (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 Pa. PG)(G)) P.e (G)Y), PG) Pie (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 the test environments. We provide formal definitions and examples of these two distribution shifts in Appendix A Here, inspired by the prior study 112, we offer formal definition to measure the graph covariate shift.
Definition 2.1 (Graph Covariate Shift) Let Ptr and Pte denote the probability functions of the training and test distributions. We measure the covariate shift between distributions Pe. and Pte
1 f (Pr(()) - Pio
GCS( Pro)
(1)
2
where S= {g Pu) Pa)(g)= 0}. which covers the features (e.g., environmental features) that do not overlap between the two distributions.
03.
GCS(.-) always bounded in [0, 1]. The issue of graph covariate shift 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, majority of the strategies [7] 5, I8 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 shared input space across training and test data. This assumption,
operate under the assumption however, often fails on covariate shift. It presents substantial challenges for these models when it
3

---

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 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 & and Ere, they follow distributions Pir and Pte, and satisfy: GCS(Pc.. Pie) > where c (0, 1) represents the degree of covariate shift. We aim to use the data collected from training environments Eu,, and learn a powerful graph classifier f G Y that performs well in all possible test environments Eroi
arg min sup ET((((())))))
ecE.e
where v)(( is the empirical risk on the environment and E((,-) is the loss function.
3
Methodology
To solve Problem 2.2 our idea to generate new graphs through data augmentation. In this section, first propose two principles for graph augmentation. Guided by these principles, design novel graph augmentation method, AIA, which can effectively address the covariate shift issue.
3.1 Two Principles for Graph Augmentation
We 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 func- tion 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 distri- bution P. From the perspective of data instances, it emphasizes the discrepancy in the environments between the generated graphs and the original graphs. Since does not give constraints on stable features, we here propose the second principle for graph augmentation.
Principle 3.2 (Stable Feature Consistency) Given set of graphs {g} with a corresponding stable feature set {geta = (Asta X.ta)). ). Let Let T-) denote an augmentation function that augments graphs (T())) with a corresponding stable feature set {gsta (Asta, Xata). Then T(.) should meet A[] sta
] + 0 and X[ xt 0, where |1 is the Frobenius norm.
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 origi- nal 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 G Rd is a graph encoder to yield d-dimensional representations, and +(-) Rd Y is a classifier. To comply with Principle 3.1 we need do OOD exploration. Inspired by distributionally robust optimization
30 31, we consider the following optimization objective:
min
sup sumpLEBlK(((o..)))))......
where P and 1 are the original and explored data distributions, respectively. -(..)) is distance met- ric between two probability distributions. The solution to Equation (3) guarantees the generalization within a robust radius p of the distribution P. To better measure the distance between distributions,

---

Stable Feature Generator
Stable Features
Augmented Graphs
GNN Classifier
Input Graph
Training
Environmental Features
Adversarial Augmenter
Figure 2: The overview of Adversarial Invariant Augmentation (AIA) Framework.
as suggested by 32 we adopt the Wasserstein distance
as the distance metric. The distance metric function can be defined as D(P, P) inf per(FP E. ( g) where T(P, P) the set of
33
34
all couplings of P and P: .(. is the cost function. Studies [35] 34 also suggest that the distances in representation space typically correspond to semantic distances. Hence, define the cost function in the representation space and give the transportation cost as c(T g) = (T) - h(gg 11.. denotes the "cost" of augmenting the graph g to g. We observe that it is difficult to set a proper p. Instead, we consider the Lagrangian relaxation for fixed penalty coefficient ry. Inspired by [2), we can reformulate Equation 3 follows:
min smmpBB//((000000-0002)))-----
(4)
where (b((g)))) sup ( (T) -7(()g)))))))
g)). And we define (gf y) as the robust surro- gate loss. If we conduct gradient descent on the robust surrogate loss, we will have:
Vo - ),)); where T
Vol(((T)
arg max { (T) vc(..9))
(5)
ges
I" is an augmented view of the original data g. Hence, to achieve OOD exploration, we need to perform graph data augmentation via Equation (5) the original data g.
3.3 Implementations of AIA
Equation (5) endows the ability of OOD exploration to data augmentation, which makes the aug- mented data meet the discrepancy principle. 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 depicted 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 To, (-) and stable feature generator (-). which generate masks for nodes and edges of graphs. They have the same structure and are parameterized by 0 and 02 respectively. Given an input graph g = (A, X) with n nodes, mask generation network first obtains the node representa- tions via GNN encoder (..). To judge the importance nodes and edges, adopts two MLP networks MLP1(.) and MLP2(-) to generate the soft node mask matrix M Raal and edge mask matrix M Rn for graph data, respectively. In summary, the mask generation network can be decomposed as:
Z = kg))).
M. = (MLP MG (MLP)((2,,
(6)
where Z Randd is node representation matrix, whose i-th row 2; Z[i.:] denotes the representation of node i, and (-) is the sigmoid function that maps the mask values M and M.₃ to [0, 1].
5

---

Adversarial Invariant Augmentation. To estimate 3" in Equation 5 we define the following
adversarial learning objective:
max (Ladv = Er. (T) yo(To, (g).
(7)
01
Then we can augments the graph by (g) (A o M ad X o M where o is the broadcasted
element-wis wise product. Although adversarially augmented graphs guarantee environmental discrep- ancy, it might destroys the stable parts. Therefore, we utilize the stable feature generator To, (.) to capture stable features and combine them with different environmental features. Following the sufficiency and invariance conditions [17 objective as:
3
18
5
8 29),
we define the stable feature learning
min {Lsta Rr, (T(((T + 6,82
(8)
where ğ (A M X M*) " is the augmented graph. It adopts the mask combination strategy: M (1 M M ady + M and M (1" - M - @ M ady M where M and M sta are generated by To, (-). 1 and 1s are all-one matrices, and if there is no edge between node and node j, then we set 1 to 0. Now we explain this combination strategy. Taking M as an example, since M denotes the captured stable regions via (.). 1 M represents the complementary parts, which are environmental regions. M represents the adversarial perturbation. so (1" - M @ Mady is equivalent to applying the adversarial perturbation on environmental features, meanwhile, sheltering the stable features. Finally, +Msta 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 rM, k A) = (E,, / - A) (Ei, [[M 0J/k- - A). where & is the total number of elements be constrained, 11 {0, 1} is an indicator function. The first term penalizes the average ratio close to x. while the second term encourages an uneven distribution. Given a graph with nodes and m edges, we define the regularization term for adversarial augmentation and stable feature learning as:
Gregi = [r, (Mad n, Aa) r rM,, [Mdvv m, Aa)]
(9)
G EP.,, [rMM n, A.) (M m, m. Aa)],
(10)
where X₂ (0. is the ratio of stable features, we usually set A for adversarial learning, which can alleviate excessive perturbations. The algorithm is provided in Appendix D.I
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 optimization objective and the discrepancy principle. Recalling the optimization objective of Equation B we aspire to identify a distribution P that can manifest within a Wasserstein ball [36]. which is centered on distribution P. with distance P serving the radius. Under appropriate conditions, we find that our learning optimization objective can establish a close connection with OOD exploration.
Proposition 4.1 Consider probubility distribution P defined over measurable space (L,F), where Q2 denotes the sample space and F a o o-algebra on Q. We construct two Wasserstein balls with P at their center and radii p1 and P2 respectively. Utilizing Equation 3 generate two distinct distributions, P1 and Pz, within the space (R,F). (f) P is an isotropic distribution; (ū) 30,, 22 Ω such that P Pi (x= x1) P2 <02); ; (iii) P1 P2, then have GCS(P, s GCS(P, Pb).
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

---

Table 1: Performance on synthetic and real-world datasets. Numbers in bold indicate the best performance, while the underlined numbers indicate the second best performance.
Motif
CMNIST
Molbbbp
Molhiv
Type
Method
base
color
seaffold
size
scaffold
size
ERM
68.66.4.25
51.74.2.88
28. 60.1.7
68. 18.10
78. 78.29.376
69.58.251
39,94.2.7
General
IRM
70.65a411
51.41
27.83.2.13
67.22.t.15
77.56.248
67.97.1.54
59.00.2.2
Generalization
GRoupDRO
68.24.8.2
51.95
29.07.3.14
66.4742.39
79.274243
70.64.2.57
58, 28.98.2.16
VREx
71.47.5.9
52.67a5.54
28. 28.48.2.87
68. 1.00
78.76.2.7
70.77.284
58.53 58.53,288
DIR
62.07.x.75
52.27
63.20
66 86.
76. 40.40
68.07
58.08+231
CAL
65.6324.29
51.18
27.99
68 28.06.2.50
79. 19.50.4.11
67.37.361
57.95.224
GSAT
62.80.11441
53.20
28. 18.1741.
66. 66.78.1445
75.63.983
68.66+1335
58, 58.06+188
Graph
Gencralization
OOD-GNN
61.10=787
=2.61
26.49.2.94
66, 1..732
79.48.4.19
70,46:199
60.60s327
StableGNN
57.07=14.10
46.93=8.55
28,38a3.49
66, 1.33
77.47:4.9
68 18.3
56,71
CIGA
66.43211.31
49. 49.14
32.22.2.67
64.92.2.9
65. 65.9813.31
69.40.2.39
59,55x2.6
DisC
51.084305
50.39
24.99
67.12.2.11
16.59.10.09
68.07.1.75
58. 58.76.091
Ddge
45.08=444
45.6344.1
22.65.2.90
66.49.
78.32+944
70.78+1.5
58.53,1.0
GREA
56.74.921
14.13:
29.02.3.26
(9.722 1.66
77.34:3.52
67.79.2.5
60.71+220
Graph
FLAG
61.1225.99
51.6624.14
32.30:2.09
67. 69.69
79.26
68.45.2.30
60,59:2.5
Augmentation
M- M-Mixup
70,08=3.2
51 51.482491
26.47.3.45
68. 68.75.0.34
78.92.244
68 68.88.265
59.03-3.11
G Mixup
59.66u7.13
52.81
31.85.5.22
67 1.64
78.5544.16
70.01
59,34.243
AIA (ours)
73.64.515
55.85.7.98
36.37.4.44
70.79.1.3
81.03.5.15
71.15.1.1
61.64.337
5
Experiments
In this section, we conduct extensive experiments to answer the following Research Questions:
RQI: Compared to existing efforts, how does AIA perform under covariate shift?
Can the proposed AIA achieve the principles of environmental feature discrepancy and stable feature consistency, thereby alleviating the graph covariate shift?
RQ2:
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 seaffold 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 [.]
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
CAL
S
GSAT
38
OOD-GNN 39
StableGNNN
[4).
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 and have the following observations:
Most generalization and augmentation methods fail under covariate shift. VREx achieves 2.8 % 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), (1..,111.991000000 11.9% and 123.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 signifi- cant improvements. For two types of covariate shifts on Motif, AIA surpasses ERM by 4.98% and 4.1 1.11% %, respectively. In contrast to the large performance variances on different datasets achieved by
7

---

Original graphs
Augmented graphs
Figure 3: Visualizations of the augmented graphs via AIA.
Table 2: Covariate shift comparisons with different augmentation strategies.
Motif (base)
Motif (size)
CMNIST (color)
Molbbbp (scaffold)
Method
Aug-Train
Aug-Test
Aug-Train
Aug-Test
Oua-Train
Aug-Test
Aug-Train
Aug-Test
Original
0
0.55720.411
0.52220.21
0
0.490za
0
0.41920099
Ddge
0.772.0213
0.515
0.851
0.161.0271
0.627+0.186
0.539-036
0.758.0
0.737
FLAG
4.001
0,533
0.002aam1
0.507
0.003au.022
0.442an002
0.001
0.41320008
G-Mixup
0.690a0116
0.472.0043
0.816
0,299.0.44
0.408
0.351
0,551
0.545
ATA (ours)
0.369.0.19
0.462.0063
0.649
0.098
0.516x0.106
0.307=n.10s
0.422
0.39340.08
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. 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 this work, we also carefully check the performance of AIA under correlation shift, and the results are presented in Appendix B
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 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 Pang with the training Pur or test distribution Pt.. "Aug-Train" and "Aug-Test" represent GCS( P. Pagg Pur) and GCS( (aang, Pte), we make these observations.
respectively. From the results Table 2
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, 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- Aug-Train and Aug-Test. G-Mixup notably augments Aug- Truin by generating OOD samples, but doesn't '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 environmental feature discrepancy. Simultancously, the imposed consistency constraint 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 intu- ition that augmentations with greater degrees of freedom yield better performance 42]. Ac- cordingly, propose conditional entropy to
Table 3: Augmentation Diversity.
Method
Full Graph
Env. Feature
Sta. Feature
DropEdge
0 0.99920065
0.933=0.229
0.971 0.259_o.m 1016
AIA (ours)
4.56
0.508-n
8

---

(basc)
Molocoo (ccafood)
75
75
Peformance
70-
Meo Sia
70
70-
ROVA
63-
65-
65-
60-
60
60
55
Hobboo
0.7
0.01
0.5
1.5 5.0
feature
Adversarial perally
Figure 4: (Left): Different components in AIA. (Middle): Different ratios X of stable features. (Righty: Different penalties r. Dashed lines denote the ERM.
measure the diversity of generated data: H(G])) EEclE, [[ where & 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 B 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 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,
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-mase ooft-mask values. These results lead us several noteworthy observations.
Visualization Analyses. AIA primarily perturbs the environmental features, while leaving the stable components undisturbed. In the Motif dataset, the base-gre aasc-graph represents ladder and the motif- graph 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 augmen- tation (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 suboptimal outcomes. Conversely, the sole use of ad 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 explore the sensitivities of ratio X, and penalty coefficient 7. The results are displayed in Figure4[[Midaa well on Motif and Molbbbp, while Molhiv is better in 0. 1-0.3. It indicates that X, is dataset-sensitive hyper-parameter that needs careful tuning. For 7, the appropriate values range from 0.151.5.
and (Right). X with 0.3~0.8 performs
6 Related Work
enlarges the training distribution by perturbing features 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 [4 45
Graph Data Augmentation 22
23 43
in graphs. Recent studies [44

---

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 paradigm for handling distribution shifts on graphs. The pioneering works K 17 leverage the causal invariance principle model the invariant predictive patterns in data for the OOD generalization purpose. With the similar spirit, GREA [88] 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.z.. molecular graphs g and recommender systems [
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]
39
40
[] 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 5]]. The latter task is called Out-of-Distriutton Detection in the literature and serves as another under-explored area that has different technical aspect from the present work. Due space constraints, we put more discussions of other related studies in Appendix G
17
38
48
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 graph augmentation method, AIA, designed to bolster the academic community 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 (92270101 14, U19A2079, 62302321) and the University Synergy Innovation Program of Anhui Province (GXXT- 2022-040). This work is also sponsored by Ant Group through CCF-Ant Research Fund and CCF. AFSG Research Fund.
References
[1] Zheyan Shen, Jiashuo Liu, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards dusto--disttiuution generalization: A survey, arXiv preprint arXiv:2108.13622 2021.
[2] Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. Good: A graph out-of-distribution benchmark. arXiv preprint arXiv:2206.0844 2022.
[3] Qitian Wu. Hengrui Zhang, Junchi Yan, and David Wipf. Handling distribution shifts graphs: An invariance perspective. In ICIR, 2022.
[4] Haoyang Li, Xin Wang. Ziwei Zhang, and Wenwu Zhu. Out-of-distributio generalization on graphs: A survey. arXiv preprint arXiv: 0772722.7777... 2022.
[5] 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.
10

---

[6] Chenxiao Yang. Qitian Wu, Jiahua Wang. and Junchi Yan. Graph neural networks are inherently good generalizers: Insights by bridging GNNs and MLPs. ICLR, 2023.
[7] Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. G-mixup: Graph data augmentation for graph classification. In ICML, pages
2022.
[8] Nianzu Yang, Kaipeng Zeng. Qitian Wu, Xiaosong Jia, and Junchi Yan. Learning substructure invariance for out-of distiiibbtttion molecular representations. In NeurlPS, 2022.
[9] Chenxiao Yang, Qitian Wu, Qingsong Wen, Zhiqiang Zhou, Liang Sun, and Junchi Yan. Towards out-of-distribution sequential event prediction: causal treatment. In NewrIPS, 2022.
[10] 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.
[11] Yuan Gao, Xiang Wang, Xiangman He, Zhenguang Liu, Huamin Feng, and Yongdong Zhang. Alleviating structural distribution shift in graph anomaly detection. In WSDM, pages 357-365, 2023.
[12] 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 -iittiiiittitttion generalization. In CVPR, pages 7947-7958, 2022.
[13] Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre Alvise-Rebutie Ira Ktena, Taylan Cemgil, et al. A fine-grained analysis on distribution shift. In ICLR, 2022.
[14] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, arXiv preprint arXiv: (tt.:::900.0999/ 2019.
[15] Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In ICML, pages 1448-1458, 2020.
[16] 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.
[17] Yingxin Wu, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. Discovering invariant rationales for graph neural networks. In ICLR, 2022.
[18] Gang Liu, Tong Zhao, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Graph rationalization with environment-based augmentations. In KDD, pages 1069-1078, 2022.
[19] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, (()):113-530 530, 2018.
[20] 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.
[21] Shaohua Fan, Xiao Wang, Yanhu Mo, Chuan Shi, and Jian Tang. Debiasing graph neural networks via learning disentangled causal substructure. In NeurIPS, 2022.
[22] Kaize Ding, Zhe Xu, Hanghang Tong, and Huan Liu. Data augmentation for deep graph learning: A survey, arXiv preprint vrXiv:2202.08235 2022.
[23] Tong Zhao, Gang Liu, Stephan Günnemann, and Meng Jiang. Graph data augmentation for graph machine learning: A survey, arXiv preprint arXiv:2202.0887/ 2022.
[24] 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.
[25] Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolutional networks on node classification. In ICLR, 2020.
11

---

[26] Yiwei Wang, Wei Wang, Yuxuan Liang, Yujun Cai, and Bryan Hooi. Mixup for node and graph classification. WWW, pages 3663-3674, 2021.
[27] Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. In NeurIPS, 2020.
[28] Susheel Suresh, Pan Li, Cong Hao, and Jennifer Neville. Adversarial graph augmentation to improve graph contrastive learning. In NeurlPS, pages 15920- 15933, 2021.
[29] Haoyang Li, Ziwei Zhang, Xin Wang, and Wenwu Zhu. Learning invariant graph representations for nuttof-diittibuutinn generalization. In NewrIPS, 2022.
[30] Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review. arXiv preprint arXiv: 1908.. 05659, 2019.
[31] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In ICLR, 2020.
[32] Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. In ICLR, 2018.
[33] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In ICML, pages 214 223, 2017.
[34] 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.
[35] Alexey Dosovitskiy and Thomas Brox. Generating images with perceptual similarity metrics based on deep networks. NIPS, 2016.
[36] Jacho Lee and Maxim Raginsky, Minimax statistical learning with wasserstein distances. NeurlPS, 31, 2018.
[37] David Krueger, Ethan Caballero, feem-Henrrk Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang. Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapo- lation (rex). In ICML, pages 5815-5826, 2021.
[38] Siqi Miao, Miaoyuan Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic attention mechanism. In ICML, pages 15524-15543, 2022.
[39] Haoyang Li, Xin Wang, Ziwei Zhang, and Wenwu Zhu. Ood-gnn: Out-of-distribution general- ized graph neural network. IEEE TKDE, 2022.
[40] Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. Generalizing graph neural networks out-of-distribution graphs arXiv preprint arXiv:2111. 2021.
[41] 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-of- distribution generalization on graphs. In NeurlPS, 2022.
[42] Raphael Gontijo-Lopes, Sylvia Smullin, Ekin Dogus Cubuk, and Ethan Dyer. Tradeoffs in data augmentation: An empirical study. ICLR, 2021.
[43] Jaemin Yoo, Sooyeon Shim, and Kang. agdeliag augmentation for accurate graph classification. In WWW, pages 1281-1291, 2022.
[44] Mucong Ding, Kezhi Kong, Jiuhai Chen, John Kirchenbauer, Micah Goldblum, David Wipf, Furong Huang, and Tom Goldstein. closer look at distribution shifts and out-of-distributtio generalization on graphs. In NeurIPSW, 2021.
[45] Chaochao Lu, Yuhuai Wu, José Miguel Hermández-Lobato, and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. In ICLR, 2021.
12

---

[46] Qi Zhu, Natalia Ponomareva, Jiawei Han, and Bryan Perozzi. Shift-robust gnns: Overcoming the limitations of localized graph training data. NeurIPS, 34:27965-279777 2021.
[47] 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.
[48] Junchi Yu, Jian Liang, and Ran He. Mind the label shift of augmentation-bassd graph ood generalization. In CVPR, 2023.
[49] Zenan Li, Qitian Wu, Fan Nie, and Junchi Yan. Graphde: A generative framework for debiased learning and out-of-distribution detection graphs. In NeurIPS, 2022.
[50] Qitian Wu, Yiting Chen, Chenxiao Yang, and Junchi Yan. Energy-based out-of-distribution detection for graph neural networks. In ICLR, 2023.
[51] Yuxin Guo, Cheng Yang. Yuluo Chen, Jixi Liu, Chuan Shi, and Junping Du A data-centric framework to endow graph neural networks with oul-o--distribution detection ability, In KDD, 2023.
[52] 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.
[53] Zheyan Shen, Peng Cui, Kun Kuang, Bo Li, and Peixuan Chen. Causally regularized learning with agnostic data selection bias. In ACM MM, pages 411 11914419 2018.
[54] 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.
[55] Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. NeurIPS, 33:1 1996-12007 2020.
[56] Yue He, Xinwei Shen, Renzhe Xu, Tong Zhang, Yong Jiang, Wenchao Zou, and Peng Cui. Covariate-shift generalization via random sample weighting. AAAI, 2023.
[57] Emanuel Parzen. On estimation of probability density function and mode. The annals of mathematical statistics, 136)),10 1962.
(58) 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.
[59] 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.
[60] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In ICLR, 2019.
[61] Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR, 2017.
[62] Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. In ICMI., pages 1725-1735. PMLR, 2020.
[63] Kiran K Thekumparampil. Chong Wang, Sewoong Oh, and Li-Jia Jia Li. Attention-basee graph neural network for semi-supervised learning. arXiv preprint arXiv: 1803.0373, 2018.
[64] 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.
[65] Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In ICML, pages 6814-6814. PMLR, 2021.
13

---

[66] Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful in an -istriidittiitttinn generalization problem? arXiv preprint arX1v:2008.01883, 2020.
[67] Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In ICML, pages 145-155. PMLR, 2020.
[68] 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.
[69] 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.
[70] 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-113) 2018.
[71] Robert Geirhos, Jörn- Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shorteut learning deep neural networks. Nature Machine Intelligence, 2(1 2(11):665 (1:666666733 663,, 2020.
[72] 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 generaliza- tion. IEEE TKDE, 2022.
[73] Jinheon Back, Dong Bok Lee, and Sung Hwang. Learning extrapolate knowledge: Transductive few-shot out-of-graphe link prediction. NeurIPS, 33:546-560, 2020.
[74] Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. In ICML, pages 7313-7324. PMLR, 2021.
[75] Massimiliano Mancini, Zeynep Akata, Elisa Ricci, and Barbara Caputo. Towards recognizing unseen categories in unseen domains. In ECCV, pages 466-483. Springer, 2020.
[76] Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In ICML, pages 10-18. PMLR, 2013.
[77] Yuning You, Tianlong Chen, Zhangyang Wang, and Yang Shen. Graph domain adaptation via theory-grounded spectral regularization. ICLR, 2023.
[78] Davide Buffelli, Pietro Lio, and Fabio Vandin. Sizeshiftreg: a regularization method for improving sezereeneralittionn in graph neural networks. In NeurIPS, 2022.
[79] 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.
[80] 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.
[81] Junfeng Fang, Xiang Wang, An Zhang, Zemin Liu, Xiangnan He, and Tat-Seng Chua. Coopera- tive explanations of graph neural networks. In WSDM, pages 316-624, 2023.
[82] 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.
[83] 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.
[84] 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.
14

---

[85] 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-300) 3024, 2022.
[86] 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
2023.
[87] 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):17461, 2023.
[88] 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.067. 2022.
15

---

A Correlation Shift and Covariate Shift
Following graph data generation process 17 29
can observe that environmental features easily change outside the training distribution, owing to their unstable nature. Hence, distribution shifts mainly caused by the environmental features [2 13 2] Specifically, we define the joint distribution of training and test data as Par (G, Y) and Pio (G, Y), respectively. Since their joint distribution can be rewritten Prr Y) Pir (Y)G)) (G) and
PY (G), we can find that there exist two main reasons that lead to the distribution shift + We give intuitive examples in Figure 5 and formal definitions of these two distribution shifts.
Correlation shift Par GY|G) + Pto(Y (G) = PG)(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 112 concept shift (2]. Formally, correlation shift describes the conditional distribution GY # (Y
(Y]G) If there exist environmental features in the test distribution that the model has not seen during training, it will also result in 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 (G) Pre (G). In Definition 2.1. we also quantitatively measure the covariate shift between (G) and (G).
Covariate shift (G) + Pc
is worth noting that within the computer vision domain, the general covariate shift is frequently synonymous with sample selection bias [52) [3] [54] [6]. 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 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 our future work.
Stable feature
Training data
Training data
Vi House
Yg: Cycle
Environmental feature
Test data
Test data
Tree Intermmeeeenn conoonanieee
Ladder
Data G Ceneration Mocbanism
Correlation Shift: Prr(Y|)) * (Y
Covartate Shift: Pc.(G) = P()(G)
Figure 5: Intuitive examples of correlation shift and covariate shift
B Proofs
In this section, we provide the detailed proofs to our proposition. We start with the definition of the Wasseerstei distance, Dw (P,,, P, P'), between two distributions and pt:
Dw Dw P')= inf Se-pee
|| yll dd. y)
meT(P,P')
where T(P, P') represents the set of all joint distributions π =(x, y) that have P and P' as their respective marginals. Under 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 ø under P to x x' under Pl. Therefore, the infimum is attained by the coupling that deterministically transitions each point x to x x'. Consequently, the Wasserstein distance simplifies
to:
Dw OP,PPP P*) f ||0- (x x'Il dP(x)) f(p. ||=0|1
[ ||='))-dP((a)) dd(w)) =||0'||
16

---

This result establishes that the Wasserstein distance between the distributions P and P' is equivalent the magnitude of x'. Therefore, we obtain:
Dw (P,P)) S( [11 (a (+ +x))l dP)(() S/ ||2211 1 (x) =|[xi||-
Similarly, Dw Pb) = ||>|||| Consequently, if Dw (P,P)) < Dw (w P2), we can deduce that ||ا||| < l
Next, we examine the GCS measure. The covariate shift between P and Pi is given by:
GCS(P, 1 Ss. J..
PP(x) P(x: +1)) dr.
Similarly, have:
GCS(P, - Js. IP(x) P(æ + x22) dx,
where S₁ and S2 denote the non-overlapping regions between P and Pt. and between P and Pe. respectively. Now we define another Pf that satisfies (f(x) P( + || [=2l) and let S denote
the non-overlapping regions between P and Py. Given the isotropy of PT. the integral over S is equal to the integral over S/. Therefore,
proceed to deduce:
1 P(x)) P(x x1) dx
GCS(P.P))
2 s 1
PP(x) - P( (x xi) da
2 s
GCS(P,P))
Let S {x
(||
27 [
x E Si), we can easily know that [S|] [5]=|SSp and St s S2. Thus
we have:
1
P(x) P(x + x1) dg
PP(x) - P(x + 21)l do
2 Sz
S.
1
P(x)) P(x + a2)) da
P((x) P(x + x da
2 s.
2
S
1
1
PP(x) P(x x22) dx 2 S
PP(x) - P(x x2)) dx
2 Sg
P(xx) P( + +2)) dx 2 0.
2 \o,
The integral over S2 could not be less than the integral over S₁, leading to GCS(P,P) s GCS(P, P2). Hence, complete the proof.
C Estimation of Graph Covariate Shift
In this section, we elaborate on the implementation details estimating the graph covariate shift. Without loss of generality, start with the example of estimating the graph covariate shift between the training and test distributions. Given the training set and test set D.r and Die, they follow proba- bility distribution functions Pt. and Pte. The process of estimating GCS(P,, Pte) is summarized in the following two steps:
Firstly, it is intractable to directly estimate the distribution in graph space G. Inspired by [12, can obtain the graph features and estimate the distribution feature space F. Specifically, given a sample, train binary GNN classifier f to distinguish which distribution it comes from, where s()) + e h, h(-)= G is a graph encoder, and ((.) adopt the pre-trained GNN encoder h to extract graph features.
{0, 1} is binary classifier. Then
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) approximate the integrals in Definition 2.1
17

---

We summarize these implementations in Algorithm 1 In lines 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 Dtr and test dataset Di:: Batch size N; Loss function 8: GNN / 4 h; Importance sampling size M; Threshold e.
Ensure: Estimated covariate shift GCS(Pur Pte)
Initialize parameters of
2: # Train graph classifier
while not converge do
4: Sample batch Btr + (gi, (:)) liet Dtr and relabel all + 0
5: Sample batch Bi + {9:. ) ] Due and relabel all y:
6: B Bt Bie
7: for each (gs. 3/4 do
8: Compute loss ((f(g:) ys) and proppgaeee gradients
9: end for
Update the parameters of 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ür and Fie via encoder h
14: F For Fif
15: Scale F to zero mean and unit variance
16: @ + fit by KDE the distribution of FF
17: Split F to recover the original partition Flr
18: Pu.,
fit by KDE the distributions of Fy Fu Fz
19: # Estimate the covariate shift
Initialize GCS(Pts. Pie) 0
21: fort + (1,... M) do
2 sample from @
if (.) < c or (z) < then
24:
GCS(Pur Pue) + GCS(Pu. Poo) (P)(()) - (1)())1/((2))
25: end if
26: end for
27:
GCS(Pur, P.o) P.
Algorithm 2: Adversarial Invariant Augmentation
Require:
Training Di:: Adversarial augmenter To, (-): Stable feature generator Top (+);
GNN classifier f(.) with parameters 0; Learning rates o, B; Batch size N; Stable
feature ratio Ag; Penalty ?.
1: Randomly initilize 6, 01. 0g
2: while not converge do
3: Sample batch Bi, +- (gt.
Drr
4: for each (gr, 3/) Btr do
5: M M Madv Te, (g:) // adversarial perturbations sta + (gi // regions of stable features
6: M... M
7: M - (MMM... @ M adv M sta // augment edges
8:
M (1 M M adv M // augment nodes
9:
Ti + (A, o M X, @ M*) // augmented graph
10:
end for
11:
Compute Ladv - Lreg, via Equation (7) and (9)
12:
Compute Leta + Lreg, via Equation (8) and (10)
Update parameters of adversarial augmenter via gradient ascent:
13:
01 O1 aVo, (Lady Lrog,)
14:
Update parameters GNN and stable feature generator via gradient descent:
0 - 0 ( (Laaa fregg): 0g 02 BV 9, (Lata freg,)
15: end while
18

---

Table 4: Statistics of graph classification datasets.
Dataset
Motif
CMNIST
Molhbbp
Malhiv
Covariate shift
base
size
color
scaffold
size
scaffold
size
Graph#
18000
18000
42000
1631
1633
24682
26169
Truin
Avg. node##
17.07
16.93
75,00
22.49
27.02
26.25
27.87
Avg. edge#
48.89
43.57
1392.76
48.43
58.71
56.68
60.20
Graphit 3000
3000
7000
204
203
2773
Val
Avg. node#
15.82
39.22
75.00
33,20
12.06
24.95
15.55
Avg. edge#
33.00
107.03
1393.73
71.84
24.27
54.53
32.77
Graphit 3000
3000
7000
204
203
4108
3961
Test
Avg. node#
14.97
87.18
75.00
27.51
12.26
19.76
12.09
Avg. edge#
31.54
239.65
1393.60
59.75
24.87
40.58
24.87
Class#
3
10
2
2
Table 5: Hyper-parameter details of AIA.
Dataset
Motif
CMNIST
Molbbbp
Molhiv
Covariate shift
base
size
color
scaffold
size
scaffold
size
Backbone (layer-hidden)
4-300
4-300
4-300
4-64
4-32
4-128
4-128
Augmenter (layer-hidden)
2-300
2-300
2-300
2-64
2-32
2-128
2-128
Generator -hidern
2-300
2-300
2-300
2-64
2-32
2-128
2-128
Optimizer Learning rate o
Adam
Adam
Adam
Adam
Adam
Adam
Adam
le-3
le-3
le-3
le-3
Se-3
le-3
le-2
Learning 3
5c-3
5e-3
Sc-3
lc-3
Se-3
1c-2
1c-2
Stable feature ratio Xs
0.5
0.5
0.5
0.5
0.5
0.1
0.1
Adversarial penalty
0.2
0.2
0.2
0.5
0.5
0.5
0.5
D Implementation Details
D.1 Algorithm
We summarize the detailed implementations of AIA in Algorithm 2 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.
We alternately optimize the
D.2 Datasets
In this paper, we conduct experiments on graph OOD datasets 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, they be regarded environmental features. The statistics of the datasets are summarized in Table 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 composed of a aase-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 the base- gase-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 tech- niques (59] We define color the environmental features create the covariate shift. Specifically, we color digits with 7 different colors, where five of them are adopted for training while the re- maining two are used for validation and testing.
Molbbbp & Molhiv: These are molecular datasets collected from MoleculeNet 19 the scaffold and graph size (i.e., node number) as the environmental features create two types of covariate shifts. For scaffold shift, we follow 2 and use scaffold split to create training, validation
We define
19

---

Table 6: Performance over diverse backbones.
Motff
Molbbbp
Backbone
Method
scaffold
EEM
67.41a4.47
50.76.292
66.44.181
47.7
M-Mixup
68.83-404
51.50.495
67.09.0.57
78.42+2.71
GON
FLAG
59.8745.61
50,99.415
66.03.2.59
78. 78.7642.54
AIA (ours)
72.39.5.77
54.87.5m
69.13.176
80.5314.43
EEM
67.8 77
11.74
67.64
77.96.4.2
M- M-Mixup
67.53-433
52.31-515
66.6440.50
76.67+2.69
GONI
FLAG
58.6745.8
50.18.441
66.72.2.52
78.55
(ours)
72.70=414
53.23:625
67.93:149
79.72.337
FEM
66.53_4.42
51.16.8m
66.51.1.77
77.41+181
M-M M-Mixup
69.25x3.9
51.374525
66.95.0443
77.21
GAT
FLAG
59.53,5.56
$1.32,445
67,36.2.45
77.8742.1
(ours)
71.95=x.
54.38:0.3
69.21:142
78.49
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
Metries
We adopt classification accuracy as the metric for Motif and CMNIST. As suggested by 20 we ROC-AUC for Molhiv and Molbbbp datasets. In addition, we use GCS(P, Q) 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 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, adopt GIN 660 as the default architecture to conduct all experiments. We tune the hyper-parameters in the following ranges: a and 3 ε (0.01, 0.005, 0.001 X2 ε {0.1,.... 0.9): Y (0.01,0.1,0.) 0.2, 0.5, 1.0, 1.5, 3.0, 5.0). The hyper-parameters are summarized in Table 5
D.5 Baseline Settings
For more comprehensive comparison, we selected 16 baselines. In this section, we give 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
CAL [) GSAT (38) DropEdge
25 GREA 18 FLAG 24,, G-Mixup 7. CIGA
[1 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 not publicly available. We reproduce them based on the codes StableNet [16].
For RDIA 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 zero in these masks. Finally, we apply these masks to the graphs for random data augmentation. The process of stable feature learning consistent with AIA.
E More Experimental Results
E.1 Results on Correlation Shift
20

---

Original graphs
GREA
CAL
DropEdge
AIA (Ours)
Figure 6: Visualizations of captured stable features.
Table 7: Performance correlation shift.
Although this work focuses the OOD issue of covariate shift, for completeness, we also evaluate the performance of AIA under correla- tion shift. Following (2), we choose three graph OOD datasets (i.e., Motif, CMNIST, Molhiv) with three different graph features (i.e., base, color, size) create correlation shifts. For base- lines, we choose three generalization algorithms (i.e., ERM. IRM [14] VREx (37)) three graph generalization methods (i.e., DIR [17). CAL (5). OOD-GNN 33] and three data augmentation methods (i.e., DropEdge [25], FLAG [24]. M-
Method
Motif
CMNIST
Molhiv
ERM
81.44,2.4
42.87x137
63. 13.26x1.25
IRM
80.71
42.80.1.2
59. 19.99
VREx
81,56.2.14
43.
6..23x
DIR
73.25s6.37
18.78=144
66. 66.78+1.50
CAL
81.94+1.20
41.82+085
62.36.1.42
OOD-GNN
80.22.2.28
39.03.1.4
57.49.1.08
DropFdge
78.97+341
38.43.1.4
54.92,1.73
FLAG
80.91
43.41
66.44.2.3
M-Mixup
77.63
40.96.1.21
64.87.11.6
AIA (ours)
82.51+281
49.73=1.70
68.11
Mixup (6) 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 and GAT 663 for exper- iments. 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, NCII, PROTEINS, COLLAB, IMDB-B, IMDB- M. For training settings, we follow CAL (5) and adopt GIN (60] as our backbone model. The experimental results shown in Table For the results, we can observe that method can achieve the best performance over different datasets.
E.4 More Visualizations
To demonstrate the superiority of our method, 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 easily observe that our method 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 and m, respectively. Let N denote the batch size, and I, denote the numbers layers in the GNN backbone, adversarial augmenter and stable feature generator, respectively. d, dg and de the
21

---

Table 8: Performance comparisons on TU datasets.
Method
MUTAG
NCII
PROTEINS
COLLAB
IMDB-B
IMDB-M
ERM
89.4227.40
82.71
76.21
82.08
73,40=3.76
51.53=297
CAL
89.91=834
83.89
76.92+3.1
82.68a125
74.13=5.21
52.60x236
DropEdge
1=.41
82.35
74,40
80.59+214
72.34-5.83
51.0623.44
FLAG
89.45=7.20
82.67=2.12
76.89+3.66
82,48
73.37=4.9
12.162
M M-Mixup
89.83=7.67
83.8922.38
76.76+3.10
82.90+1.43
74.07=4.76
52.89
AIA (ours)
90.34:7.75
84.12:264
77.92:3772
82.98
74.23=8.10
53.02.2.76
Table 9: Running time, model size and performance improvement.
FRM
CAT.
DIR
(curs)
Dutaset
Running Time
Model Size
Running Time
Model
Performance Improvemeet
Running Time
Medel Size
Performance Improvement
Running
Model
Performance
Size
Time
Size
Improvement
Moui
0Kh S1mm 19s
1.515M
01h 57m 15s
2.2133
1 2.949%
0th 50mm 27s
2.158M
5.005
0th 49tm DNs
2.366M
7.535%
CMNEST
0th 56en 92x
1.517M
07th 28m 49s
2:244M
1 2.15%
02h 34m 35s
3.16M
16.0695
02h- 4Sm 30s
2.57MM
1 272.17%
Molbbhp
00h
11m 58s
1.515M
00h
18m 22s
2.213M
1 0.81%
00h 20m 1
2.158M
2. 2.12%
00h 19t 376
2.366M
3.722%
Molhhv
000
27mh:
196
1.515M
000 46m 14s
2.213M
1 8.239
000 5kas
40s
2.158M
2.709%
00h 55m 46s
2.366M
1225455
dimensions of hidden layers the GNN backbone, adversarial augmenter and stable feature generator, respectively.
Time Complexity. The time complexity of the adversarial learning objective is ONN 2/md)). For the stable feature learning objective, the time complexity is O(N (L,mde + 22md)). For the regularization terms, the time complexity is ((2N m)). ). For simplicity, we assume la le and do de. Hence, the time complexity of forward propagation is O22
+ +2mmd+ n m)).
Model Size. In addition to the GNN backbone model, also introduce two small networks for adversarial augmentation and stable feature learning. In our implementations, the parameters of ATA around twice as large those of the original GNN model. The running time and model size shown in Table
G More Related Studies
OOD Generalization I 65. 66 has been widely explored. Inspired by invariant learning and causal theory, a series of general algorithms 67 68,
66 70) have recently been proposed to solve the 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 statistical correlation that is inconsistent in training and test data. Thus, the models prefer learn spurious correlations and rely 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
69
OOD problem. Recent studies 122 2 D3
contexts, such as node classification 46 3 47 graph classification [39
38 4%
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] [8], explainability [81] [7 38, 82). and computational efficiency [83] 84 85). Concurrently, effoots are being made to extend these principles to diverse applications, such as
graphs [8] recommender systems 9 10 and anomaly detection
molccular graphs
86 87].
Comprehensive Comparisons with EERM 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 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. contrast, we explicitly model the environmental and stable features. Hence, we can effectively identify stable and environmental features and explicitly
22

---

Table 10: Comparisons with EERM.
EERM
AIA
Is it specifically designed
Scope
for covariate shift?
x
Can uuvirommental/tabli
Separubility
X
features be sepurated?
Can environmental Teatures be identified explicitly? model
Environmental Feature Discrepancy
Mask model To, (-)
environmental features?
Metrie for envimnmental Di screpancy
GCS( P. P)
Generation principle for environmental features
"Blindly" maximize V. [.[R(e)]
Maximize GCS(P,
Can leatures
be identified explicitly?
x
Stable Feature Consistency
How to model
Mask model T-,()
stable features?
Learning principles for
ming V.[R(e)]
Sufficieeey/lnseeeedddeee
stable features
Theoretical basis
IRM
DRO Robust P D s
Generalization
Generalization scope
K environments
separate them from data. Thirdly, we also design metric, GCS(P,P) P). which can effectively measure the discrepancy of the environmental features for our augmented data. And 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 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 p. We summarize the above detailed discussions in Table T0)
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 given training set, such as perturbing original graph node features or graph structures. Hence, 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 advanced data augmentation methods, such 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 backhone GNN. This training strategy may make the training process unstable, so the performance of AIA may experience 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 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.
23


Reviews:
[{"review_id": "txUDNqofrS", "reviewer": "Reviewer_reTY", "summary": "This 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!", "questions": "see weaknesses", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ 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.", "weaknesses": "1. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for rating our paper as \"well written and easy to follow\", \"clearly\", \"sufficiently novel\", and \"well presented\"! According to your concerns, we provide the following responses. \n\n### Q1: Concerns about adversarial augmenter, stable feature generator, masks and regularizer.\n\n- **Difference between stable and adversarial mask.** We utilize two distinct modules to create stable and adversarial masks. The stable mask aims to highlight the stable features in the graph data. Ideally, we'd like the mask value for the stable feature to be 1, while the other parts are 0. The adversarial mask, on the other hand, is used to modify the data. With this mask, we aim to disrupt the graph features, such as by removing certain nodes or edges, in order to create out-of-distribution data.\n\n- **Preventing mask overlap or disjoint.** We understand your concern about the potential for conflict between the two masks. To address this, we use a strategy to combine the masks in a way that the adversarial perturbation doesn't harm the stable features. This is explained in lines 207-218 of the original paper. Basically, the stable mask ($M_{sta}$) highlights the stable regions. Anything not covered by this mask, represented as $1-M_{sta}$, is the complementary part. The adversarial mask ($M_{adv}$) represents the adversarial perturbation. By applying the operation $(1-M_{sta})\\odot M_{adv} + M_{sta}$, we apply the adversarial perturbation to the complementary parts \"$(1-M_{sta})\\odot M_{adv}$\", while preserving the stable features \"$+ M_{sta}$\" . This ensures that the final augmented data includes both new environmental features and original stable features.\n\n- **Purpose of regularization term.** The regularization term is used to guide the masks to closely match our expectations (i.e., the mask value should be near 0 or 1). This prevents the masks from converging to a trivial solution during training, thus helping to avoid overlap or conflict between the masks. For the stable mask, we aim for it to highlight the stable feature, and for its value to be close to 0 or 1. We set the stable ratio as $\\lambda_s$ and design two regularization terms to enforce this. For the adversarial mask, we limit its ratio to 1 to prevent it from creating excessive perturbations, like deleting all components to optimize the adversarial goal.\n\n### Q2. How does this approach handle size generalization?\n\nGreat question! Current studies, especially invariant learning, suggest that size distribution shift occurs due to (size) differences in environmental features during training and testing. Let's take molecular graphs as an example. The size of scaffolds, such as carbon chains or rings, can greatly vary. Our method aims to maintain stable features while modifying environmental features to create new graphs. During training, the adversarial mask might alter the graph size by removing varying numbers of nodes or edges (as seen in Figure 2 or Figure 3). This way, the model can handle graphs with different environmental feature sizes, reducing sensitivity to graph size. Therefore, the model can generalize to unseen graphs of different sizes during testing. We appreciate your suggestion and will include this explanation in the final version for clarity.\n\n### Q3. More insights into the results. For example, the bigger benefits are observed in the Motif (base) setting.\n\nYou've made a great observation. Our method shows significant improvements on Motif (base) mainly because it experiences larger covariate shift. We've quantified this covariate shift (GCS) for different datasets (as shown in Table 2). The GCS of Motif (base) is 0.557, which is larger than others. As our method is specifically designed to tackle covariate shift, and other baseline methods have limitations in this regard, our improvement is more noticeable on Motif (base). Further, upon reviewing our results in light of your suggestion, we've added two additional insights:\n\n- Our method outperforms GREA and attains the best performance in terms of size generalization, even though GREA, another data augmentation method, also performs well (second place on motif and molhiv). GREA uses two types of coarse-grained augmentation, namely, environment replacement and removal, which might produce new graphs of varying sizes. In contrast, our method applies a more detailed environmental feature augmentation, specifically by removing certain nodes or edges. Therefore, it considers a broader size scope than GREA, which further improves the performance.\n- Invariant learning methods like DIR, CAL, and DisC struggle to perform well as they find it challenging to generate new environmental features. However, covariate shift is very evident in these datasets (see Table 2). By creating new environments during data augmentation, our method ensures better generalization under covariate shift.\n\nThank you for your suggestions. We will include these insights in the final version of our paper.\n\n### Q4. Will the benefits persist with a graph transformer?\n\nIndeed, they do. We conduct experiments on Molbbbp using four different graph transformers [1-4] as shown **in our one-page PDF** (Table 1). The results show that our method continues to provide benefits across these diverse graph transformers.\n\n[1] Do Transformers Really Perform Bad for Graph Representation? NeurIPS 2021\n\n[2] Representing Long-Range Context for Graph Neural Networks with Global Attention, NeurIPS 2021\n\n[3] Rethinking Graph Transformers with Spectral Attention, NeurIPS 2021\n\n[4] Recipe for a General, Powerful, Scalable Graph Transformer, NeurIPS 2022"}, {"review_id": "qI8hIIFlQN", "reviewer": "Reviewer_bjaC", "summary": "The 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.", "questions": "Please answer my concerns in the previous section.\n", "limitations": "No concern.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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. ", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and comments! We have responded to your concerns as follows. We sincerely hope that addressing these concerns will help change your rating of our work!\n\n### Q1.  Covariate shift is a well-studied problem.\n\n- **We're talking about Graph Covariate Shift.** We agree with you that covariate shift is well-studied in general settings [1-4]. But our paper looks at a different and new problem called \"graph covariate shift.\" This is a new area in graph learning and hasn't been studied a lot, as Reviewer ExMB said, \"*It is an important OOD area of graph learning*\". When we say \"covariate shift is a less explored type of OOD,\" we mean this specific graph learning context. We are very sorry if our words seem to be about the general settings.  And we will thoroughly fix all unclear parts in our paper and make sure it's clear.\n- **Graph OOD is not the same as general OOD.** Traditional OOD problems usually deal with simple tasks like computer vision. For these tasks, inputs are things like variables or image features. But graphs are more complex. This means that graph OOD problems also have to deal with structure distribution shifts, not just shifts in the features. Therefore, it is non-trivial to study covariate shift in graph learning tasks (comments from Reviewer ExMB). \n\nThanks for your suggestion! If our paper is accepted, we promise to add more discussions about covariate shift, such as [1-4].\n\n### Q2. The connection to covariate shift is misleading. & The paper assumes the \"common patterns\" should exist in both training and test distribution.\n\n### Q3. How the stable feature generator can find stable features in the dataset?\n\nYour concerns (Q2 and Q3) might come from a misunderstanding about \"stable features\" and the invariance assumption. We first begin with the following two clarifications:\n\n- **Clarification 1: The Invariance Assumption.**  OOD generalization is challenging and even impossible without any assumptions [5]. Hence, the invariance assumption has been introduced by [5] and applied by follow-up works [6-11] as a cornerstone assumption for data generation which enables reasonable problem-solving and analysis for OOD problems. Specifically, stable feature (or invariant feature) causally determines the label and their relationship is invariant across distributions.  Let's use the example of a cow in different backgrounds (grassy or sandy). The cow object is the \"stable feature\" since the relation from the cow object to the label is invariant across environments (various backgrounds). If we train our model (like an image classifier) on images with a grassy background, it should learn to focus on the cow (stable feature) and ignore the background. This way, it can work well even when we show images with new environments (e.g., sandy background). Similarly, stable features also widely exist in the graph domain (e.g., molecular functional groups [9-11]).\n- **Clarification 2: Stable Features Are Not Always Common Patterns.** The reviewer seems to think that stable features are always common patterns. We didn't make this claim in our paper. This misunderstanding could come from Motif dataset, where the stable feature looks similar a lot. It is a synthetic dataset that are commonly-used in existing studies. They [10, 11] usually use Motif to intuitively check how well a method works. In our work, we only assume the causal mechanisms are invariant [5-11]. We don't say that the stable feature itself always looks the same. For example, in the cow picture, the cow could be different colors or shapes. In molecular graphs (e.g., Molhiv and Molbbbp in our paper), different parts might cause a certain property. Thanks, we will emphasize this in our final paper.\n\nNow let's move on to your two questions:\n\n**Answer to Q2:** In our paper (LINE-110), we provide a definition for covariate shift. Our work aims to address graph covariate shift based on the invariance assumption. This assumption for OOD has been explored in existing studies, such as in Ood-Bench [8] and GOOD [9]. Ood-Bench proves that dealing with covariate shift is challenging and even impossible without any assumptions about data generation. GOOD creates many graph-domain datasets that show graph covariate shift, also based on the invariance assumption. We use GOOD's datasets and therefore follow their settings and assumptions.\n\n**Answer to Q3**: From the above clarifications, we are not looking for \"common pattern\", but looking for \"stable features\". Stable features are substructures of the data that can causally determine the label, and their relationship is invariant across environments. The model can make right predictions based on stable features. Now let's check our loss: $\\ell(f(T_{\\theta_2}(g)), y) + \\ell(f(\\widetilde{g}), y)$. The first item \"$\\ell(f(T_{\\theta_2}(g)), y)$\" requires the model to make correct predictions based on the estimated stable features, and $\\ell(f(\\widetilde{g}), y)$ requires invariant predictions across different environments. We also confirmed through intuitive visualizations (see Figure 3) and experiments (see Table 3) that we can approximately find stable features.\n\nReferences:\n\n[1] A Theoretical Analysis on Independence-driven Importance Weighting for Covariate-shift Generalization, ICML 2022\n\n[2] Causally Regularized Learning with Agnostic Data Selection Bias, ACM MM 2018\n\n[3] Rethinking Importance Weighting for Deep Learning under Distribution Shift, NeurIPS 2020\n\n[4] Covariate-Shift Generalization via Random Sample Weighting, AAAI 2023\n\n[5] Invariant models for causal transfer learning, JMLR, 2018.\n\n[6] Invariant risk minimization, ArXiv, 2019.\n\n[7] OOD generalization via risk extrapolation, ICML, 2021.\n\n[8] Quantifying and Understanding Two Dimensions of OOD Generalization, CVPR 2022\n\n[9] GOOD: A Graph OOD Benchmark, NeurIPS 2022\n\n[10] Discovering Invariant Rationales for GNNs, ICLR 2022\n\n[11] Causal Attention for Interpretable and Generalizable Graph Classification, KDD 2022"}, {"review_id": "RgC8rWWDx9", "reviewer": "Reviewer_ExMB", "summary": "The 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.", "questions": "- What is the GNN backbone used in AIA for the experiments?\n- How will the choice of GNN backbones affect the final outcome?", "limitations": "The paper addresses the limitations in the appendix.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We gratefully thank you for the positive feedback and constructive comments! To address your concerns, we provide point-to-point responses, and some necessary results are in our **one-page PDF** in the Global Response above.\n\n### Q1. Adversarial training is inefficient.\n\nIn Appendix F (from Supplementary Material), we have already discussed our optimization complexity and model size. In our training steps, a single step of updating the parameters of the adversarial generator can achieve good performance, instead of using a multi-step approach similar to PGD. So our complexity is acceptable. ***In our one-page PDF (Table 5)***, we count the running time and model size of ERM, DIR [1], CAL [2] and our AIA on NVIDIA 3090 (24GB GPU). Our running time is about 1.5-2 times that of the base model. With two additional small networks, our model size is roughly 1.5 times the size of the original model. And our method is comparable to the current SOTA methods DIR, CAL in terms of running time and model size. Hence, we believe we achieve a better performance-complexity trade-off considering that ours can achieve significant performance improvements. In practical applications, we believe that these additional complexities are acceptable. Of course, we accept your suggestion and will put these discussions and results in our final version. And we will also find ways to reduce our complexity in our future work.\n\n### Q2. What is the GNN backbone used in AIA for the experiments?\n\nFor a fair comparison, we uniformly choose GIN [3] as the backbone for all algorithms.\n\n### Q3. How will the choice of GNN backbones affect the final outcome?\n\n***In our one-page PDF (Table 2)***, we also selected three different backbone models (GCN [4], GCNII [5] and GAT [6]) for experiments to answer your question. In addition, we have also conducted experiments based on diverse graph transformer backbones ***in our one-page PDF (Table 1)***. Note that, our observations and conclusions remain the same with the new results. If our paper is accepted, we promise to add these new results to the final version of the paper.\n\n[1] Discovering Invariant Rationales for Graph Neural Networks, ICLR 2022\n\n[2] Causal Attention for Interpretable and Generalizable Graph Classification, KDD 2022\n\n[3] How Powerful are Graph Neural Networks?, ICLR 2019\n\n[4] Semi-Supervised Classification with Graph Convolutional Networks, ICLR 2017\n\n[5] Simple and deep graph convolutional networks, ICML 2020\n\n[6] Graph Attention Networks, ICLR 2018"}, {"review_id": "P5wN0pbLBM", "reviewer": "Reviewer_J9ZS", "summary": "This 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! ", "questions": "- 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? ", "limitations": "See above. Negative societal impacts, etc are adequately discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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\n", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable time and comments! We provide detailed responses as below, and some necessary results are in our **one-page PDF** in the Global Response above. We sincerely hope that addressing your concerns will help change your rating of our work!\n### Q1. Limited Novelty & Motivation are not novel & Adversarial augmentation has been explored by other methods.\nWe appreciate your thoughts but respectfully argue that our method and motivation are novel. Now we list the following aspects.\n- **Reviewer Feedback**. The novelty and motivation have been praised by other three reviewers, such as ***Reviewer ExMB***'s comment: \"The adversarial augmentation method design is novel with good intuitions and it provides new insights into the important OOD area of graph learning.\"; ***Reviewer bjaC***:\"The technique proposed by the authors is interesting, and the use of two components on graph data, is new.\"; and ***Reviewer reTY***: \"I believe the extension and its implementation to this problem is sufficiently novel.\"\n- **Our Novelty & Motivation.** Covariate shift on graphs remains largely under-explored (comments from ***Reviewer ExMB***). In this paper, we provide a clear formulation for this new problem, design a new method to solve this problem with solid theoretical analyses. Further, we also propose a new metric, to measure the graph covariate shift in diverse datasets (see Table 2 in our paper). Moreover, using invariance principle with graph augmentation for handling graph covariate shift is new and unique.\n- **Comparisons with Similar Methods (refer to our one-page PDF).** We understand your concern that our augmentation is not novel. However, most adversarial augmentation strategies [1-5] are performed on European data (e.g. image). Due to the non-Euclidean nature of graph data, it is difficult to directly transfer these strategies to graphs. Therefore, it is non-trivial to design new method in graph domain (comments from ***Reviewer bjaC***). There are few studies applying adversarial augmentation or masking to graph data. ***In our one-page PDF (Table 3)***, we list several graph augmentation methods (ADGCL [6], RGCL [7], EERM, FLAG, GREA) that are most similar to ours, and highlight the main differences. \n### Q2. Table 3 (augmentation diversity) is interesting & Compare non-random augmentation methods (e.g., FLAG).\nWe compare with non-random methods such as FLAG, G-Mixup, ADGCL, RGCL. Since ADGCL and RGCL are SSL settings, we replace our augmentation module in AIA with theirs. From the results ***in our one-page PDF (Table 4)***, FLAG mainly focuses on node feature, which limits its diversity. And FLAG does not distinguish between stable and environmental features. For G-Mixup, despite having high environmental diversity, struggles to maintain stable feature consistency. Others like ADGCL and RGCL also fall short in comparison to our performance.\n### Q3. How to select hyper-parameters during tuning?\nWe provide detailed hyper-parameters in Appendix (Table 2). They are tuned within the following ranges: $\\alpha, \\beta \\in \\\\{0.01, 0.005, 0.001\\\\}$; $\\lambda_s \\in \\\\{0.1, ..., 0.9\\\\}$; $\\gamma \\in \\\\{0.01, 0.1, 0.2, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0\\\\}$. The best parameters are chosen based on accuracy on the validation set. For basic settings like batch size, optimizer, we keep them consistent with baseline models.\n### Q4. Complexity & Adversarial objective and masking operation is expensive. & Can you add some run-times/model sizes to the appendix for later drafts?\nIn Appendix F, we have already discussed our time complexity and model size. Here we discuss more details about complexity, model size and running time. We would also be happy to include these details in our final version if the paper is accepted.\n- **Adversarial objective complexity .** In our training steps, a single step of updating the parameters of the adversarial generator can achieve good performance, instead of using a multi-step approach similar to PGD. So our complexity is acceptable.\n- **Masking operation complexity.**  Rather than computing the scores between all pairs of nodes (i.e., $O(n^2)$). We predict the score for each edge. Therefore, our time complexity is $O(m)$. For node feature, the time complexity is $O(n)$, so the total complexity for masking operation is $O(n+m)$. Furthermore, we focus on graph classification. In OGB dataset, the largest dataset has an average of only 243 nodes and 2261 edges. Hence, our approach is suitable to handle them and will not include more time complexity.\n- **Running time and Model size.** ***In our one-page PDF (Table 5)***, we count the running time and model size of ERM, DIR, CAL and our AIA on NVIDIA 3090. Our running time is about 1.5-2 times that of the base model. With two additional small networks, our model size is roughly 1.5 times the size of the original model. And our method is comparable to the current SOTA methods DIR, CAL in terms of running time and model size. Hence, we believe we achieve a better performance-complexity trade-off considering that ours can achieve significant accuracy gains. In practical applications, we believe that these additional complexities are acceptable. Of course, **we accept your suggestion and will put these discussions and results in our paper**. And we will also find ways to reduce our complexity in future work.\n\n[1] Adversarial AutoAugment, ICLR 2020 \n\n[2] Maximum-Entropy Adversarial Data Augmentation for Improved Generalization and Robustness, NeurIPS 2020\n\n[3] AugMax: Adversarial Composition of Random Augmentations for Robust Training, NeurIPS 2021\n\n[4] Rethinking the Effect of Data Augmentation in Adversarial Contrastive Learning, ICLR 2023\n\n[5] Harnessing OOD Examples via Augmenting Content and Style, ICLR 2023\n\n[6] Adversarial Graph Augmentation to Improve Graph Contrastive Learning, NeurIPS 2021\n\n[7] Let Invariant Rationale Discovery Inspire Graph Contrastive Learning, ICML 2022"}]

Author rebuttal:
# Global Response and One-page PDF from Authors 

We appreciate all the reviewers' efforts for reviewing this submission. We are delighted that our paper was noted for being "*well written and easy to follow*", "*clear",* "*sufficiently novel*" by **Reviewer reTY**; "*interesting*", "*new*" by **Reviewer bjaC**; offering "*new insights*", "*novel with good intuitions*" by **Reviewer ExMB**; and "*well-motivated*" by **Reviewer J9ZS**. To address the reviewers' concerns, we have included new experimental results and a method comparison table in our ***one-page PDF*** below. Here, we summarize the main points raised by the reviewers and our responses.

- **[@Reviewer bjaC], Covariate shift & Invariance Assumption.** We understand that your main concern stems from a misunderstanding of our proposed "stable feature". We have provided a thorough explanation for you and clarified the relationship between our work and the invariance assumption.  We apologize for any confusion caused by vague phrases such as "covariate shift is a less explored type of OOD", and will thoroughly revise our paper. Finally, we sincerely hope to engage in further discussions to address your concerns and improve your rating of our work.

- **[@Reviewer J9ZS], Novelty & Motivation.** We understand that you question the novelty and motivation of our graph augmentation method. However, other reviewers have strongly recognized the novelty of our approach. Transferring general augmentation methods to graph data is challenging due to its non-Euclidean nature. We have also compared our approach with other similar methods in the graph domain (***Table 3 in PDF***). For your interest in our augmentation diversity experiments, we have included the results in ***Table 4 in PDF***. If our paper is accepted, we'd be glad to include all the running time and model size results (***Table 5 in PDF***) in our appendix. We sincerely hope that addressing these concerns will help change your rating of our work.


- **[@Reviewer J9ZS, ExMB], Complexity & Running Time & Model Size.** Thanks to the modest size of the two extra networks we implemented and our one-step adversarial training, our approach maintains manageable complexity and model size parameters. For the running time and model size comparisons (***Table 5 in PDF***), it is evident that our method achieves a preferable performance-complexity balance compared to other baselines.

- **[@Reviewer reTY, ExMB], Different Backbones.** We conducted experiments using a broader range of backbone models (***Table 1 and 2 in PDF***), including GIN, GCN, GCNII, GAT, and various Graph transformers. The results confirm that our findings and conclusions remain consistent across different backbones.

Meta-review: This paper received four reviews, of which three are quite positive. There have been a lot of discussions during rebuttal and major concerns have been addressed.
Decision: Accept (poster)