SlowGuess/NeurIPS_2023_OCR huggingface dataset:
DrugOOD datasets. To evaluate the OOD performance in realistic scenarios with realistic distribution shifts, we also include three datasets from DrugOOD benchmark [28]. DrugOOD is a systematic OOD benchmark for AI-aided drug discovery, focusing on the task of drug target binding affinity prediction for both macromolecule (protein target) and small-molecule (drug compound). The molecule data and the notations are curated from realistic ChEMBL database [54]. Complicated distribution shifts can happen on different assays, scaffolds and molecule sizes. In particular, we select DrugOOD-lbap-core-ec50-assay, DrugOOD-lbap-core-ec50-scaffold, DrugOOD-lbap-core-ki-assay, DrugOOD-lbap-core-ki-scaffold, and DrugOOD-lbap-core-ki-size, from the task of Ligand Based Affinity Prediction which uses ic50 measurement type and contains core level annotation noises. We directly use the data files provided by the authors. For more details, we refer interested readers to Ji et al. [28].
G.2 Baselines and Evaluation Setup
During the experiments, we do not tune the hyperparameters exhaustively while following the common recipes for optimizing GNNs. Details are as follows.
GNN encoder. For a fair comparison, we use the same GNN architecture as graph encoders for all methods. By default, we use 3-layer GIN [75] with Batch Normalization [27] between layers and JK residual connections at the last layer [74]. The hidden dimension is set to 32 for Two-piece graphs, CMNIST-sp, and 128 for SST2, and DrugOOD datasets. The pooling is by default a mean function over all nodes. The only exception is DrugOOD datasets, where we follow the backbone used in the paper [28], i.e., 4-layer GIN with sum readout.
Interpretable GNN backbone. As mentioned in Sec. 2 that most of the existing invariant graph learning approaches adopt the interpretable GNN as the basic backbone model for the whole predictor  $f = f_{c} \circ g$ , where  $g: \mathcal{G} \to \mathcal{G}_{c}$  is a featurizer GNN and  $f_{c}: \mathcal{G}_{c} \to \mathcal{V}$  is a classifier GNN.  $g$  first calculates the sampling weights as in  $\widehat{G}_c$  for each edge. More formally, given a graph  $G$  containing  $n$  nodes, a soft mask is predicted through the following equation:
$$
Z = \operatorname {G N N} (G) \in \mathbb {R} ^ {n \times h}, M = \mathrm {a} (Z, A) \in \mathbb {R} ^ {n \times n},
$$
where  $a$  calculates the sampling weights for each edge using a MLP:  $M_{ij} = \mathrm{MLP}([Z_i,Z_j])$ . Based on the continuous sampling score  $M$ ,  $g$  could sample discrete edges according to the predicted scores [55]. For two-piece graph datasets and DrugOOD datasets, we will directly use the score to reweight the messaging passing process along the edge, as we empirically find it yields more stable performance. While for CMNIST-sp and Graph-SST2, we will sample a ratio  $r\%$  of all edges for each graph. The ratios adopted are  $80\%$  and  $60\%$ , respectively, following previous works [8, 28]. Meanwhile, to improve the stability of the subgraph extractor, we adopt a layernorm [66] following the practice of [55].
Besides, we also have various implementation options for obtaining the features in  $\widehat{G}_c$ , for further obtaining  $h_{\widehat{G}_c}$ , as well as for obtaining predictions based on  $\widehat{G}_s$ . Following previous works [55], we will adopt the same GNN encoder for the two GNNs in the interpretable GNN backbone, and feed the raw graph inputs to the classifier GNN. The contrastive loss is obtained via the graph representations of the sampled subgraph by the classifier GNN. For classifying  $G$  based on  $\widehat{G}_s$ , we use a separate MLP downstream classifier in the classifier GNN  $f_c$ .
Optimization and model selection. By default, we use Adam optimizer [34] with a learning rate of  $1e - 3$  and a batch size of 128 for all models at all datasets. Except for CMNIST-sp, we use a batch size of 256 to facilitate the evaluation following previous works [55]. To avoid underfitting, we pre-train models for 20 epochs for all datasets by default. While in two-piece graphs, we find pre-training by 100 epochs yields more stable performance. To avoid overfitting, we also employ an early stopping of 5 epochs according to the validation performance. Meanwhile, dropout is also adopted for some datasets. Specifically, we use a dropout rate of 0.5 for all of the realistic graph datasets, following previous works [8, 28].
<sup>4</sup>https://drugood.github.io/
The final model is selected according to the performance at the validation set. All experiments are repeated with 5 different random seeds of  $\{1,2,3,4,5\}$ . The mean and standard deviation are reported from the 5 runs.
Implementations of Euclidean OOD methods. When implementing IRM [3], V-Rex [38] and IB-IRM [2], we refer the implementations from DomainBed [23]. Since the environment information is not available, we perform random partitions on the training data to obtain two equally large environments for these objectives following previous works [8, 12]. Moreover, we select the weights for the corresponding regularization from  $\{0.01, 0.1, 1, 10, 100\}$  for these objectives according to the validation performances of IRM and stick to it for others, since we empirically observe that they perform similarly with respect to the regularization weight choice. For EIIL [13], we use the author-released implementations about assigning different samples the weights for being put in each environment and calculating the IRM loss.
Implementations of invariant graph learning methods. We implement GSAT [55], GREA [45], CAL [64], MoleOOD [77], GIL [41], DisC [17], and CIGA [8], according to the author provided codes (if available).
- GREA [45]: We use a penalty weight of 1 for GREA as we empirically it does not affect the performance by changing to different weights.
- Interpretable ratio: same as others;
- Penalty weight: 1;
Number of environments: N/A;
- GSAT [55]: We follow the recommendations of the released implementations by the authors.
- Interpretable ratio:  $70\%$ ;
- Penalty weight: 1;
- Decay ratio:  $10\%$ ;
- Decay interval: pretrain epoch//2;
Number of environments: N/A;
- CAL [64]: We follow the recommendations of the released implementations by the authors.
- Interpretable ratio: same as others;
- Penalty weight:  $\{0.1, 0.5, 1.0\}$ ;
Number of environments: N/A;
- MoleOOD [77]: We tune the penalty weights of MoleOOD with values from  $\{1e - 2, 1e - 1, 1, 10\}$  but did not observe much performance differences. Hence we stick the penalty weight as 1 for all datasets.
- Interpretable ratio: N/A;
- Penalty weight: 1;
- Number of environments: same as others;
GIL [41]: We follow the recommendations of the paper.
- Interpretable ratio: same as others;
- Penalty weight:  $\{1e - 5, 1e - 3, 1e - 1\}$ ;
Number of environments: same as others;
- DisC [17]: We tune only the  $q$  weight from  $\{0.9, 0.7, 0.5\}$  in the GCE loss as we did not observe performance differences by changing the weight of the other terms.
- Interpretable ratio: same as others;
-  $q$  weight:  $\{0.9, 0.7, 0.5\}$ ;
Number of environments: same as others;
- CIGA [8]: We follow the recommendations of the released implementations by the authors..
- Interpretable ratio: same as others;
- Penalty weight:  $\{0.5, 1, 2, 4, 8, 16, 32\}$ ;
- Number of environments: N/A;
GALA:
- Interpretable ratio: same as others;
- Penalty weight:  $\{0.5, 1, 2, 4, 8, 16, 32\}$ ;
- Environment assistant: {vanilla GNN, XGNN};
- Sampling proxy: {label predictions, cluster predictions};
- Number of environments: same as others;
All of the graph learning methods adopt an interpretable GNN as the backbone by default. The only exception is MoleOOD, we follow the original implementation while using a shared GNN encoder for the variational losses to ensure the fairness of comparison. Besides, for DisC, we find the soft masking implementation in two-piece graphs will incur a severe performance degeneration hence we use a ratio of  $25\%$  for the interpretable GNN backbone.
For environment inferring methods, we search the number of environments
- Two-piece graphs: fixed as 3 (since there are 3 spurious graphs);
- CMNIST-sp: 2 (since there are 2 environments);
Graph-SST2:  $\{2,3,4\}$  following previous practice [41];
- DrugOOD datasets:  $\{2,3,5,10,20\}$  following previous practice [77].;
Implementations of GALA. For a fair comparison, GALA uses the same GNN architecture for GNN encoders as the baseline methods. By default, we fix the temperature to be 1 in the contrastive loss, and merely search the penalty weight of the contrastive loss from  $\{0.5, 1, 2, 4, 8, 16, 32\}$  according to the validation performances, following the CIGA implementations [8]. By default, we implement the environment assistant as a ERM model, and adopt directly the environment assistant predictions to sample possible and negative graph pairs. Nevertheless, as discussed in Sec. 4 that there could be multiple implementation choices for the environment assistant and the use of its predictions. We hence also try with XGNN based environment assistant model and clustering based proxy predictions. By default, the selection of the environment assistant model is performed via best training performance, as which encourages a better fit to the dominant subgraph patterns, while we also try the model selection with best validation performance in DrugOOD datasets and find it empirically sometimes leads to better performance. All the options for the selection of the environment assistant models depend on the validation performance. For Two piece graphs, EC50-Scaffold, EC50-Size, Ki-Assay, Ki-Scaffold, CMNIST-sp and Graph-SST2, we find implementing the environment assistant as a ERM model already yield impressive improvements. While for the other DrugOOD datasets, we implement the environment assistant as an interpretable GNN trained with ERM and cluster the learned graph representations of the model to sample positive and negative pairs.
Since GALA imposes a strong regularization to the data that may hinder the learning of graph representations, we pre-train the model by 10 epochs using ERM and then impose the GALA penalty implemented as one-side contrastive loss as discussed in Sec. F. When the numbers of positive and negative pairs are extremely imbalanced, we will upsample the minor groups by a factor of  $\{2,3,4\}$ , depending on the validation performance.
G.3 Software and Hardware
We implement our methods with PyTorch [59] and PyTorch Geometric [18]. We ran our experiments on Linux Servers installed with V100 graphics cards and CUDA 10.2.
G.4 Computational analysis
Table 6: Averaged total training time of different methods.
Datasets	Two-piece graphs	EC50-Assay	CMNIST-sp	Graph-SST2
ERM	435.85±2.14	80.45±10.27	315.84±5.55	374.31±1.28
XGNN	673.82±0.81	126.65±17.57	591.09±11.48	722.44±48.51
GREA	1128.28±34.57	210.30±21.23	902.06±8.49	979.15±18.3114
GSAT	1205.67±62.54	228.88±25.04	791.55±15.67	949.57±97.68
DisC	1244.68±4.76	207.50±17.72	932.40±76.99	1280.77±551.97
MoleOOD	714.06±6.53	136.39±17.87	439.49±9.10	712.31±81.62
GIL	533.46±11.42	279.30±25.39	919.53±14.15	733.36±147.08
CIGA	873.49±16.21	167.63±1.10	650.94±5.01	792.10±59.12
GALA-cluster	811.41±3.20	147.97±2.05	756.41±21.63	765.32±20.86
GALA-pred	793.27±8.58	149.89±2.71	644.78±53.58	764.69±30.98
We calculate the average total training time of different methods at various datasets in seconds. As shown in Table. 6, the training of GALA (no matter with clustering based sampling or prediction based sampling) does not bring much additional overhead than its counterpart CIGA. When considering the additional training time of the assistant model with ERM, GALA costs only a competitive training time as environment generation based methods such as GREA and DisC. Notably, some methods such as DisC and GIL sometimes may be slow to converge even with the same early stop setting, which will cost even more time than the time cost by GALA plus the ERM training. Besides, the ERM training time (for a assistant model) is not much long and usually around 5mins (or 300seconds in the table).