SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- ERM [65] minimizes the empirical loss on the training set.
- IRM [19] seeks to find data representations across all environments by penalizing feature distributions that have different optimal classifiers.
- VREx [40] reduces the risk variances of training environments to achieve both covariate robustness and invariant prediction.
<sup>3</sup>https://drugood.github.io/
<sup>4</sup>http://ftp.ebi.ac.uk/pub/databases/chembl/ChEMBLdb/releases/chembl_30
Table 5: Description of different environment splits leading to molecular distribution shifts.
Environment	Explanation
Scaffold	Molecular scaffold is the fundamental structure of a molecule with desirable bioactive properties. Molecules with the same scaffold belong to the same environment. Distribution shift arises when there is a change in the molecular scaffold.
Size	The size of a molecule refers to the total number of atoms in the molecule. Molecular size is also an inherent structural characteristic of molecular graphs. The distribution shift occurs when size changes.
Assay	Assay is an experimental method as an examination or determination for molecular characteristics. Due to variations in assay environments and targets, activity values measured by different assays often differ significantly. Samples tested within the same assay belong to a single environment, while a change in the assay leads to a distribution shift.
- GroupDRO [41] minimizes the loss on the worst-performing group, subject to a constraint that ensures the loss on each group remains close.
- Coral [44] encourages feature distributions consistent by penalizing differences in the means and covariances of feature distributions for each domain.
- DANN [43] encourages features from different environments indistinguishable by adversarially training a regular classifier and a domain classifier.
- Mixup [66] augments data in training through data interpolation.
And the others are graph-specific methods:
- DIR<sup>5</sup> [29] discovers the subset of a graph as invariant rationale by conducting interventional data augmentation to create multiple distributions.
- GSAT $^{6}$  [49] proposes to build an interpretable graph learning method through the attention mechanism and inject stochasticity into the attention to select label-relevant subgraphs.
- GREA $^{7}$  [47] identifies subgraph structures called rationales by environment replacement to create virtual data points to improve generalizability and interpretability.
- CAL $^{8}$  [28] proposes a causal attention learning strategy for graph classification to encourage GNNs to exploit causal features while ignoring the shortcut paths.
- DisC $^{9}$  [32] analyzes the generalization problem of GNNs in a causal view and proposes a disentangling framework for graphs to learn causal and bias substructure.
- MoleOD $^{10}$  [25] investigates the OOD problem on molecules and designs an environment inference model and a substructure attention model to learn environment-invariant molecular substructures.
- CIGA[11] [26] proposes an information-theoretic objective to extract the desired invariant subgraphs from the lens of causality.
Experiments are conducted on one 24GB NVIDIA RTX 3090 GPU.
<sup>5</sup>https://github.com/Wuyxin/DIR-GNN
$^{6}$ https://github.com/Graph-COM/GSAT
<sup>7</sup>https://github.com/liugangcode/GREA
<sup>8</sup>https://github.com/yongduosui/CAL
$^{9}$ https://github.com/googlebaba/DisC
$^{10}$ https://github.com/yangnianzu0515/Mole0OD
11https://github.com/LFphase/CIGA
Figure 5: Hyper-parameter sensitivity analysis on the covariate-shift dataset of GOODHIV-Scaffold.
Baselines. For datasets in the GOOD benchmark, we use the results provided in the official leaderboard $^{12}$ . For datasets in the DrugOOD benchmark, we use the official benchmark code $^{13}$  to get the results for ERM, IRM, Coral, and Mixup on the latest version of datasets. The results for GroupDRO and DANN are not reported due to an error occurred while the code was running. For some baselines that do not have reported results, we implement them using public codes. All of the baselines are implemented using the GIN-Virtual [3, 35] (on GOOD) or GIN [35] (on DrugOOD) as the GNN backbone that is parameterized according to the guidance of the respective benchmark. And we conduct a grid search to select hyper-parameters for all implemented baselines.
Our method. We implement the proposed iMoLD in Pytorch [69] and PyG [70]. For all the datasets, we select hyper-parameters by ranging the code book size  $|\mathcal{C}|$  from  $\{100, 500, 1000, 4000, 10000\}$ , threshold  $\gamma$  from  $\{0.1, 0.5, 0.7, 0.9\}$ ,  $\lambda_1$  from  $\{0.001, 0.01, 0.1, 0.5\}$ ,  $\lambda_2$  from  $\{0.01, 0.1, 0.5, 1\}$ ,  $\lambda_3$  from  $\{0.01, 0.1, 0.3, 0.5, 1\}$ , and batch size from  $\{32, 64, 128, 256, 512\}$ . For datasets in DrugOOD, we also select dropout rate from  $\{0.1, 0.3, 0.5\}$ . The maximum number of epochs is set to 200 and the learning rate is set to 0.001. Please refer to Table 6 for a detailed hyper-parameter configuration of various datasets. The hyper-parameter sensitivity analysis is in Appendix C.1.
C.1 Hyper-parameter Sensitivity Analysis
Take the dataset of covariate-shift split of GOODHIV-Scaffold as an example, we conduct extensive experiments to investigate the hyper-parameter sensitivity, and the results are shown in Figure 5. We observe that the performance tends to improve first and then decrease slightly as the size of the codebook  $|\mathcal{C}|$  increases. This is because a small codebook limits the expressivity of the model, while too large one cuts the advantage of the discrete space. The effect of the threshold  $\gamma$  is insignificant and there is no remarkable trend. As the  $\lambda_{1}$ ,  $\lambda_{2}$  and  $\lambda_{3}$  increase, the performance shows a tendency to increase first and then decrease, indicating that  $\mathcal{L}_{\mathrm{inv}}$ ,  $\mathcal{L}_{\mathrm{reg}}$  and  $\mathcal{L}_{\mathrm{cmt}}$  are effective and can improve performance within a reasonable range. We also observe that the standard deviation of performance increases as  $\lambda_{1}$  increases, which may be due to the fact that too much weight on the self-supervised invariant learning objective may enhance or affect the performance. The standard deviation is the smallest when  $\lambda_{2} = 0$ , suggesting that neural networks may have more stable outcomes by learning adaptively when there are no constraints, but it is difficult to obtain higher performance. While the
$^{12}$ https://good.readthedocs.io/en/latest/leaderboard.html
<sup>13</sup>https://github.com/tencent-ailab/DrugOD
Table 6: Hyper-parameter configuration.
γ	|C|	batch-size	λ1	λ2	λ3	dropout
GOOD	HIV	scaffold	covariate	0.8	4000	128	0.01	0.5	0.1
concept	0.7	4000	256	0.01	0.5	0.1
size	covariate	0.7	4000	256	0.01	0.5	0.1
concept	0.9	4000	1024	0.01	0.5	0.1
ZINC	scaffold	covariate	0.3	4000	32	0.01	0.5	0.1
concept	0.5	4000	256	0.01	0.5	0.1
size	covariate	0.3	4000	256	0.01	0.5	0.1
concept	0.3	4000	64	0.0001	0.5	0.1
PCBA	scaffold	covariate	0.9	10000	32	0.0001	1	0.1
concept	0.9	10000	32	0.0001	1	0.1
size	covariate	0.9	10000	32	0.0001	1	0.1
concept	0.9	10000	32	0.0001	1	0.1
DrugOOD	IC50	assay	0.7	1000	128	0.001	0.5	0.1	0.5
scaffold	0.9	1000	128	0.0001	0.5	0.1	0.5
size	0.7	1000	128	0.01	0.5	0.1	0.1
EC50	assay	0.7	500	128	0.01	0.5	0.1	0.5
scaffold	0.3	500	128	0.001	0.5	0.1	0.3
size	0.7	500	128	0.001	0.5	0.1	0.1
standard deviation is the largest when  $\lambda_3 = 0$ , indicating that the commitment loss in VQ can increase the performance stability.