SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we provide the detailed proofs to our proposition. We start with the definition of the 1-Wasserstein distance,  $D_W(P, P')$ , between two distributions  $P$  and  $P'$ :
$$
D _ {W} (P, P ^ {\prime}) = \inf  _ {\pi \in \Gamma (P, P ^ {\prime})} \int_ {\mathbb {R} ^ {d} \times \mathbb {R} ^ {d}} \| \boldsymbol {x} - \boldsymbol {y} \| \cdot d \pi (\boldsymbol {x}, \boldsymbol {y}),
$$
where  $\Gamma(P, P')$  represents the set of all joint distributions  $\pi(\boldsymbol{x}, \boldsymbol{y})$  that have  $P$  and  $P'$  as their respective marginals. Under our conditions, due to  $P'(\boldsymbol{x} - \boldsymbol{x}') = P(\boldsymbol{x})$ , we can precisely determine the manner in which the mass from  $P$  was transferred to create  $P'$ . This process involves moving each point  $\boldsymbol{x}$  under  $P$  to  $\boldsymbol{x} + \boldsymbol{x}'$  under  $P'$ . Therefore, the infimum is attained by the coupling that deterministically transitions each point  $\boldsymbol{x}$  to  $\boldsymbol{x} + \boldsymbol{x}'$ . Consequently, the Wasserstein distance simplifies to:
$$
D _ {W} (P, P ^ {\prime}) = \int_ {\mathbb {R} ^ {d}} | | \boldsymbol {x} - (\boldsymbol {x} + \boldsymbol {x} ^ {\prime}) | | \cdot d P (\boldsymbol {x}) = \int_ {\mathbb {R} ^ {d}} | | \boldsymbol {x} ^ {\prime} | | \cdot d P (\boldsymbol {x}) = \| \boldsymbol {x} ^ {\prime} \|.
$$
This result establishes that the Wasserstein distance between the distributions  $P$  and  $P'$  is equivalent to the magnitude of  $x'$ . Therefore, we obtain:
$$
D _ {W} (P, \widetilde {P} _ {1}) = \int_ {\Omega} | | \boldsymbol {x} - (\boldsymbol {x} + \boldsymbol {x} _ {1}) | | \cdot d \widetilde {P} _ {1} (\boldsymbol {x}) = \int_ {\Omega} | | \boldsymbol {x} _ {1} | | \cdot d P _ {1} (\boldsymbol {x}) = \| \boldsymbol {x} _ {1} \|.
$$
Similarly,  $D_W(P, \widetilde{P}_2) = \| \pmb{x}_2 \|$ . Consequently, if  $D_W(P, \widetilde{P}_1) < D_W(P, \widetilde{P}_2)$ , we can deduce that  $\| \pmb{x}_1 \| < \| \pmb{x}_2 \|$ .
Next, we examine the GCS measure. The covariate shift between  $P$  and  $\widetilde{P}_1$  is given by:
$$
\operatorname {G C S} (P, \widetilde {P} _ {1}) = \frac {1}{2} \int_ {S _ {1}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {1}) | \cdot d \boldsymbol {x}.
$$
$$
\operatorname {G C S} (P, \widetilde {P} _ {2}) = \frac {1}{2} \int_ {S _ {2}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {2}) | \cdot d \boldsymbol {x},
$$
where  $S_{1}$  and  $S_{2}$  denote the non-overlapping regions between  $P$  and  $\widetilde{P}_1$ , and between  $P$  and  $\widetilde{P}_{2}$ , respectively. Now we define another  $\widetilde{P}_1^{\prime}$  that satisfies  $\widetilde{P}_1^\prime (\pmb {x}) = P(\pmb {x} + \| \pmb {x}_1\| \cdot \frac{\pmb{x}_2}{\|\pmb{x}_2\|})$  and let  $S_1^{\prime}$  denote the non-overlapping regions between  $P$  and  $\widetilde{P}_1^{\prime}$ . Given the isotropy of  $\widetilde{P}_1$ , the integral over  $S_{1}$  is equal to the integral over  $S_1^{\prime}$ . Therefore, we can proceed to deduce:
$$
\begin{array}{l} \operatorname {G C S} (P, \widetilde {P} _ {1}) = \frac {1}{2} \int_ {\mathcal {S} _ {1}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {1}) | \cdot d \boldsymbol {x} \\ = \frac {1}{2} \int_ {\mathcal {S} _ {1} ^ {\prime}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {1}) | \cdot d \boldsymbol {x} \\ = \operatorname {G C S} \left(P, \widetilde {P} _ {1} ^ {\prime}\right) \\ \end{array}
$$
Let  $S_2' = \{\pmb{x} + (\| \pmb{x}_2 \| - \| \pmb{x}_1 \|) \cdot \frac{\pmb{x}_2}{\|\pmb{x}_2\|} \mid \pmb{x} \in S_1'\}$ , we can easily know that  $|\mathcal{S}_1'| = |\mathcal{S}_2'|$  and  $\mathcal{S}_2' \subseteq \mathcal{S}_2$ . Thus we have:
$$
\begin{array}{l} \operatorname {G C S} (P, \widetilde {P} _ {2}) - \operatorname {G C S} (P, \widetilde {P} _ {1}) = \frac {1}{2} \int_ {\mathcal {S} _ {2}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {1}) | \cdot d g - \frac {1}{2} \int_ {\mathcal {S} _ {1}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {1}) | \cdot d \boldsymbol {x} \\ = \frac {1}{2} \int_ {\mathcal {S} _ {2}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {2}) | \cdot d \boldsymbol {x} - \frac {1}{2} \int_ {\mathcal {S} _ {1} ^ {\prime}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {1}) | \cdot d \boldsymbol {x} \\ = \frac {1}{2} \int_ {\mathcal {S} _ {2}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {2}) | \cdot d \boldsymbol {x} - \frac {1}{2} \int_ {\mathcal {S} _ {2} ^ {\prime}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {2}) | \cdot d \boldsymbol {x} \\ = \frac {1}{2} \int_ {\mathcal {S} _ {2} \backslash \mathcal {S} _ {2} ^ {\prime}} | P (\boldsymbol {x}) - P (\boldsymbol {x} + \boldsymbol {x} _ {2}) | \cdot d \boldsymbol {x} \geq 0. \\ \end{array}
$$
The integral over  $S_{2}$  could not be less than the integral over  $S_{1}$ , leading to  $\mathrm{GCS}(P, \widetilde{P}_{1}) \leq \mathrm{GCS}(P, \widetilde{P}_{2})$ . Hence, we complete the proof.
C Estimation of Graph Covariate Shift
In this section, we elaborate on the implementation details of estimating the graph covariate shift. Without loss of generality, we start with the example of estimating the graph covariate shift between the training and test distributions. Given the training set and test set  $\mathcal{D}_{\mathrm{tr}}$  and  $\mathcal{D}_{\mathrm{te}}$ , they follow probability distribution functions  $P_{\mathrm{tr}}$  and  $P_{\mathrm{te}}$ . The process of estimating  $\mathrm{GCS}(P_{\mathrm{tr}}, P_{\mathrm{te}})$  is summarized in the following two steps:
- Firstly, it is intractable to directly estimate the distribution in graph space  $\mathbb{G}$ . Inspired by [12], we can obtain the graph features and estimate the distribution in feature space  $\mathbb{F}$ . Specifically, given a sample, we train a binary GNN classifier  $f$  to distinguish which distribution it comes from, where  $f(\cdot) = \Phi \circ h$ ,  $h(\cdot): \mathbb{G} \to \mathbb{F}$  is a graph encoder, and  $\Phi(\cdot): \mathbb{F} \to \{0,1\}$  is a binary classifier. Then we can adopt the pre-trained GNN encoder  $h$  to extract graph features.
- Secondly, we prepare the features and estimate the distribution of the data via Kernel Density Estimation (KDE) [57]. Finally, we adopt the Monte Carlo Integration under importance sampling [58] to approximate the integrals in Definition 2.1.
We summarize these implementations in Algorithm 1. In lines 4 and 5, to avoid the label shift [12], we adopt sample reweighting to ensure the balance of each class.
Algorithm 1: Estimation of Graph Covariate shift
Require: Training dataset  $\mathcal{D}_{\mathrm{tr}}$  and test dataset  $\mathcal{D}_{\mathrm{te}}$ ; Batch size  $N$ ; Loss function  $\ell$ ; GNN  $f = \Phi \circ h$ ; Importance sampling size  $M$ ; Threshold  $\epsilon$ .  
Ensure: Estimated covariate shift  $\mathrm{GCS}(P_{\mathrm{tr}}, P_{\mathrm{te}})$ .  
1: Initialize parameters of  $f$   
2: # Train a graph classifier  
3: while not converge do  
4: Sample a batch  $\mathcal{B}_{\mathrm{tr}} \gets \{(g_i, y_i)\}_{i=1}^N \subset \mathcal{D}_{\mathrm{tr}}$  and relabel all  $y_i \gets 0$   
5: Sample a batch  $\mathcal{B}_{\mathrm{te}} \gets \{(g_i, y_i)\}_{i=1}^N \subset \mathcal{D}_{\mathrm{te}}$  and relabel all  $y_i \gets 1$   
6:  $\mathcal{B} \gets \mathcal{B}_{\mathrm{tr}} \cup \mathcal{B}_{\mathrm{te}}$   
7: for each  $(g_i, y_i) \in \mathcal{B}$  do  
8: Compute loss  $\ell(f(g_i), y_i)$  and back-propagate gradients  
9: end for  
10: Update the parameters of  $f$  via gradient descent and reset the gradients  
11: end while  
12: # Prepare the features for the estimation  
13: Extract training and test feature sets  $\mathcal{F}_{\mathrm{tr}}$  and  $\mathcal{F}_{\mathrm{te}}$  via encoder  $h$   
14:  $\mathcal{F} \gets \mathcal{F}_{\mathrm{tr}} \cup \mathcal{F}_{\mathrm{te}}$   
15: Scale  $\mathcal{F}$  to zero mean and unit variance  
16:  $\hat{\omega} \gets$  fit by KDE the distribution of  $\mathcal{F}$   
17: Split  $\mathcal{F}$  to recover the original partition  $\mathcal{F}_{\mathrm{tr}}'$ ,  $\mathcal{F}_{\mathrm{te}}'$   
18:  $\hat{P}_{\mathrm{tr}}, \hat{P}_{\mathrm{te}} \gets$  fit by KDE the distributions of  $\mathcal{F}_{\mathrm{tr}}', \mathcal{F}_{\mathrm{te}}'$   
19: # Estimate the covariate shift  
20: Initialize GCS  $(P_{\mathrm{tr}}, P_{\mathrm{te}}) \gets 0$   
21: for  $t \gets \{1, ..., M\}$  do  
22:  $z \gets$  sample from  $\hat{\omega}$   
23: if  $\hat{P}_{\mathrm{tr}}(z) < \epsilon$  or  $\hat{P}_{\mathrm{te}}(z) < \epsilon$  then  
24: GCS  $(P_{\mathrm{tr}}, P_{\mathrm{te}}) \gets$  GCS  $(P_{\mathrm{tr}}, P_{\mathrm{te}}) + |\hat{P}_{\mathrm{tr}}(z) - \hat{P}_{\mathrm{te}}(z)| / \hat{\omega}(z)$   
25: end if  
26: end for  
27: GCS  $(P_{\mathrm{tr}}, P_{\mathrm{te}}) \gets$  GCS  $(P_{\mathrm{tr}}, P_{\mathrm{te}})/2M$
Algorithm 2: Adversarial Invariant Augmentation
Require: Training set  $\mathcal{D}_{\mathrm{tr}}$  ; Adversarial augmenter  $T_{\theta_1}(\cdot)$  ; Stable feature generator  $T_{\theta_2}(\cdot)$  GNN classifier  $f(\cdot)$  with parameters  $\theta$  ; Learning rates  $\alpha ,\beta$  ; Batch size  $N$  ; Stable feature ratio  $\lambda_{s}$  ; Penalty  $\gamma$    
1: Randomly initialize  $\theta ,\theta_{1},\theta_{2}$    
2: while not converge do   
3: Sample a batch  $\mathcal{B}_{\mathrm{tr}}\gets \{(g_i,y_i)\}_{i = 1}^N\subset \mathcal{D}_{\mathrm{tr}}$    
4: for each  $(g_i,y_i)\in \mathcal{B}_{\mathrm{tr}}$  do   
5:  $\mathbf{M}_{\mathrm{adv}}^{a},\mathbf{M}_{\mathrm{adv}}^{x}\gets T_{\theta_{1}}(g_{i}) / /$  adversarial perturbations   
6:  $\mathbf{M}_{\mathrm{sta}}^{a},\mathbf{M}_{\mathrm{sta}}^{x}\gets T_{\theta_{2}}(g_{i}) / /$  regions of stable features   
7:  $\widetilde{\mathbf{M}}^a\gets (1 - \mathbf{M}_{\mathrm{sta}}^a)\odot \mathbf{M}_{\mathrm{adv}}^a +\mathbf{M}_{\mathrm{sta}}^a / /$  augment edges   
8:  $\widetilde{\mathbf{M}}^x\gets (1 - \mathbf{M}_{\mathrm{sta}}^x)\odot \mathbf{M}_{\mathrm{adv}}^x +\mathbf{M}_{\mathrm{sta}}^x / /$  augment nodes   
9:  $\widetilde{g}_i\gets (\mathbf{A}_i\odot \widetilde{\mathbf{M}}^a,\mathbf{X}_i\odot \widetilde{\mathbf{M}}^x) / /$  augmented graph   
10: end for   
11: Compute  $\mathcal{L}_{\mathrm{adv}} - \mathcal{L}_{\mathrm{reg}_1}$  via Equation (7) and (9)   
12: Compute  $\mathcal{L}_{\mathrm{sta}} + \mathcal{L}_{\mathrm{reg}_2}$  via Equation (8) and (10)   
13: Update parameters of adversarial augmenter via gradient ascent:   
 $\theta_{1}\leftarrow \theta_{1} + \alpha \nabla_{\theta_{1}}(\mathcal{L}_{\mathrm{adv}} - \mathcal{L}_{\mathrm{reg}_1})$    
14: Update parameters of GNN and stable feature generator via gradient descent:   
 $\theta \gets \theta -\beta \nabla_{\theta}(\mathcal{L}_{\mathrm{sta}} + \mathcal{L}_{\mathrm{reg}_2});\theta_2\gets \theta_2 - \beta \nabla_{\theta_2}(\mathcal{L}_{\mathrm{sta}} + \mathcal{L}_{\mathrm{reg}_2})$    
15: end while
Table 4: Statistics of graph classification datasets.
Dataset Covariate shift	Motif	CMNIST color	Molbbbp	Molhiv
base	size	scaffold	size	scaffold	size
Train	Graph#	18000	18000	42000	1631	1633	24682	26169
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
Val	Graph#	3000	3000	7000	204	203	4113	2773
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
Test	Graph#	3000	3000	7000	204	203	4108	3961
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	3	10	2	2	2	2
Table 5: Hyper-parameter details of AIA.
Dataset Covariate shift	Motif	CMNIST color	Molbbbnp	Molhiv
base	size	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 (layer-hidden)	2-300	2-300	2-300	2-64	2-32	2-128	2-128
Optimizer	Adam	Adam	Adam	Adam	Adam	Adam	Adam
Learning rate α	1e-3	1e-3	1e-3	1e-3	5e-3	1e-3	1e-2
Learning rate β	5e-3	5e-3	5e-3	1e-3	5e-3	1e-2	1e-2
Stable feature ratio λs	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
We summarize the detailed implementations of AIA in Algorithm 2. We alternately optimize the adversarial augmenter and stable feature generator with the backbone model, in lines 13 and 14. We adopt the learned stable features for predictions in the inference stage.
In this paper, we conduct experiments on graph OOD datasets [2] and OGB datasets [20], which include Motif, CMNIST, Molbbb and Molhiv. We follow [2] to create various covariate shifts, according to base, color, size and scaffold splitting. Base, color, size and scaffold are features of the graph data and do not determine the labels of the data, so they can be regarded as environmental features. The statistics of the datasets are summarized in Table 4. Below we give a brief introduction to each dataset.
- Motif: It is a synthetic dataset from Spurious-Motif [17, 5]. As shown in original graphs in Figure 5, each graph is composed of a base-graph (wheel, tree, ladder, star, path) and a motif (house, cycle, crane). The label is only determined by the type of motif. We create covariate shift according to the base-graph type and the graph size (i.e., node number). For base covariate shift, we adopt graphs with wheel, tree, ladder base-graphs for training, star for validation and path for testing. For size covariate shift, we use small-size of graphs for training, while the validation and the test sets include the middle- and the large-size graphs, respectively.
- CMNIST: Color MNIST dataset contains graphs transformed from MNIST via superpixel techniques [59]. We define color as the environmental features to create the covariate shift. Specifically, we color digits with 7 different colors, where five of them are adopted for training while the remaining two are used for validation and testing.
- Molbbbp & Molhiv: These are molecular datasets collected from MoleculeNet [19]. We define the scaffold and graph size (i.e., node number) as the environmental features to create two types of covariate shifts. For scaffold shift, we follow [2] and use scaffold split to create training, validation
Table 6: Performance over diverse backbones.
Backbone	Method	Motif	Molbbbp
base	size	scaffold	size
GCN	EEM	67.41±4.47	50.76±2.92	66.44±1.91	77.79±4.04
M-Mixup	68.83±4.04	51.50±4.95	67.09±0.57	78.42±2.71
FLAG	59.87±5.61	50.99±4.18	66.03±2.59	78.76±2.54
AIA (ours)	72.39±5.37	54.87±5.02	69.13±1.76	80.53±4.43
GCNII	EEM	67.84±4.74	51.74±3.15	67.64±1.84	77.96±4.02
M-Mixup	67.53±4.31	52.31±5.18	66.64±0.50	76.67±2.69
FLAG	58.67±5.88	50.18±4.41	66.72±2.52	78.55±2.52
AIA (ours)	72.70±4.14	53.23±6.25	67.93±1.69	79.72±3.37
GAT	EEM	66.53±4.42	51.16±3.22	66.51±1.77	77.41±3.81
M-Mixup	69.25±3.99	51.37±5.25	66.95±0.43	77.21±2.48
FLAG	59.53±5.56	51.32±4.48	67.36±2.45	77.87±2.31
AIA (ours)	71.95±3.39	54.38±6.32	69.21±1.62	78.49±3.20
and test sets. For size shift, we adopt the large-size of graphs for training and the smaller ones for validation and testing.
We adopt classification accuracy as the metric for Motif and CMNIST. As suggested by [20], we use ROC-AUC for Molhiv and Molbbb datasets. In addition, we use  $\mathrm{GCS}(P,Q)$  to measure the covariate shift between distributions  $P$  and  $Q$ . For all experimental results, we perform 10 random runs and report the mean and standard derivations. For augmentation diversity, we use conditional entropy to measure the diversity of generated data. We normalize it to  $[0,1]$  for better comparison. Specifically, for a given graph data, we compute the conditional entropy by collecting augmented data generated during the training process. We conduct experiments by collecting 1000 graph data, and report the mean and standard deviation.
D.4 Training Settings
We use the NVIDIA GeForce RTX 3090 (24GB GPU) to conduct all our experiments. To make a fair comparison, we adopt GIN [60] as the default architecture to conduct all experiments. We tune the hyper-parameters in the following ranges:  $\alpha$  and  $\beta \in \{0.01, 0.005, 0.001\}$ ;  $\lambda_2 \in \{0.1, \dots, 0.9\}$ ;  $\gamma \in \{0.01, 0.1, 0.2, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0\}$ . The hyper-parameters are summarized in Table 5.
D.5 Baseline Settings
For a more comprehensive comparison, we selected 16 baselines. In this section, we give a detailed introduction to the settings of these methods.
- For ERM, IRM [14], GroupDRO [31], VREx [37], and M-Mixup [26], we report the results from the study [2] by default and reproduce the missing results on Molbbpb.
- For DIR [17], CAL [5], GSAT [38], DropEdge [25], GREA [18], FLAG [24],  $\mathcal{G}$ -Mixup [7], CIGA [41] and DisC [21], they provide source codes for the implementations. We adopt default settings from their source codes and detailed hyper-parameters from their original papers.
- For OOD-GNN [39] and StableGNN [40], their source codes are not publicly available. We reproduce them based on the codes of StableNet [16].
- For RDIA in Section 5.4, it is a variant that replaces the adversarial augmentation in AIA with random augmentation. In our implementation, we use all-one matrices to create the initial node and edges masks. Then we randomly set  $20\%$  of nonzero elements to zero in these masks. Finally, we apply these masks to the graphs for random data augmentation. The process of stable feature learning is consistent with AIA.
E.1 Results on Correlation Shift
Figure 6: Visualizations of captured stable features.
Although this work focuses on the OOD issue of covariate shift, for completeness, we also evaluate the performance of AIA under correlation shift. Following [2], we choose three graph OOD datasets (i.e., Motif, CMNIST, Molhiv) with three different graph features (i.e., base, color, size) to create correlation shifts. For baselines, we choose three generalization algorithms (i.e., ERM, IRM [14], VReX [37]), three graph generalization methods (i.e., DIR [17], CAL [5], OOD-GNN [39]) and three data augmentation methods (i.e., DropEdge [25], FLAG [24], M
Table 7: Performance on correlation shift.
Method	Motif	CMNIST	Molhiv
ERM	81.44±2.54	42.87±1.37	63.26±1.25
IRM	80.71±2.81	42.80±1.62	59.90±1.17
VReX	81.56±2.14	43.31±1.03	60.23±1.60
DIR	73.25±6.37	38.78±1.45	66.78±1.50
CAL	81.94±1.20	41.82±0.85	62.36±1.42
OOD-GNN	80.22±2.28	39.03±1.24	57.49±1.08
DropEdge	78.97±3.41	38.43±1.94	54.92±1.73
FLAG	80.91±1.04	43.41±1.38	66.44±2.32
M-Mixup	77.63±1.12	40.96±1.21	64.87±1.36
AIA (ours)	82.51±2.81	49.73±1.70	68.11±1.82
Mixup [26]). The experimental results are shown in Table 7. We can observe that AIA can also effectively alleviate the correlation shift. These results demonstrate that AIA learns better stable features by encouraging environmental discrepancy, which can effectively break spurious correlations that are hidden in the training data.
E.2 Results on Diverse Backbones
We select three different GNN backbone models (GCN [61], GCNII [62] and GAT [63]) for experiments. From the results in Table 6, our observations and conclusions remain the same with the diverse backbones.
E.3 Results on More Real-world Datasets
To demonstrate the effectiveness of the proposed AIA, we also conduct experiments on commonly used TU datasets [64], which include MUTAG, NCI1, PROTEINS, COLLAB, IMDB-B, IMDB-M. For training settings, we follow CAL [5] and adopt GIN [60] as our backbone model. The experimental results are shown in Table 8. For the results, we can observe that our method can achieve the best performance over different datasets.
E.4 More Visualizations
To demonstrate the superiority of our method, we also visualize the captured stable features by AIA and compare them with other baselines. The results are displayed in Figure 6. From the results, we can easily observe that our method can find stable parts more accurately than other baseline methods.
F Complexity Analyses
Firstly, we define the average numbers of nodes and edges per graph in the dataset to be  $n$  and  $m$ , respectively. Let  $N$  denote the batch size,  $l$ ,  $l_{a}$  and  $l_{c}$  denote the numbers of layers in the GNN backbone, adversarial augmenter and stable feature generator, respectively.  $d$ ,  $d_{a}$  and  $d_{c}$  are the
Table 8: Performance comparisons on TU datasets.
Method	MUTAG	NCI1	PROTEINS	COLLAB	IMDB-B	IMDB-M
ERM	89.42±7.40	82.71±1.52	76.21±3.83	82.08±1.51	73.40±3.78	51.53±2.97
CAL	89.91±8.34	83.89±1.93	76.92±3.31	82.68±1.25	74.13±5.21	52.60±2.36
DropEdge	86.11±9.41	82.35±3.77	74.40±3.10	80.59±2.14	72.34±5.83	51.06±3.04
FLAG	89.45±7.20	82.67±2.12	76.89±3.66	82.48±1.79	73.37±4.94	52.16±2.70
M-Mixup	89.83±7.67	83.89±2.38	76.76±3.40	82.90±1.43	74.07±4.76	52.89±2.84
AIA (ours)	90.34±7.75	84.12±2.64	77.92±3.72	82.98±1.76	74.23±5.10	53.02±2.76
Table 9: Running time, model size and performance improvement.
Dataset	ERM	CAL	DIR	AIA (ours)
Running Time	Model Size	Running Time	Model Size	Performance Improvement	Running Time	Model Size	Performance Improvement	Running Time	Model Size	Performance Improvement
Motif	00h 51m 19s	1.515M	01h 37m 15s	2.213M	↓ 2.98%	01h 50m 27s	2.158M	↓ 5.03%	01h 49m 08s	2.366M	↑ 7.55%
CMNIST	01h 56m 32s	1.517M	02h 28m 49s	2.244M	↓ 2.13%	02h 34m 35s	2.161M	↑ 16.08%	02h 45m 30s	2.373M	↑ 27.17%
Molbbp	00h 11m 58s	1.515M	00h 18m 22s	2.213M	↑ 0.81%	00h 20m 11s	2.158M	↓ 2.12%	00h 19m 37s	2.366M	↑ 3.72%
Molivh	00h 27m 19s	1.515M	00h 46m 14s	2.213M	↓ 3.23%	00h 58m 40s	2.158M	↓ 2.59%	00h 55m 46s	2.366M	↑ 2.54%
dimensions of hidden layers in the GNN backbone, adversarial augmenter and stable feature generator, respectively.
Time Complexity. The time complexity of the adversarial learning objective is  $\mathcal{O}(N(l_a md_a + 2lmd))$ . For the stable feature learning objective, the time complexity is  $\mathcal{O}(N(l_c md_c + 2lmd))$ . For the regularization terms, the time complexity is  $\mathcal{O}(2N(n + m))$ . For simplicity, we assume  $l_a = l_c$  and  $d_a = d_c$ . Hence, the time complexity of a forward propagation is  $\mathcal{O}(2N(l_a md_a + 2lmd + n + m))$ .
Model Size. In addition to the GNN backbone model, we also introduce two small networks for adversarial augmentation and stable feature learning. In our implementations, the parameters of AIA are around twice as large as those of the original GNN model. The running time and model size are shown in Table 9.
G More Related Studies