SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} _ {\text {p r e d}} = \sum_ {i = 1} ^ {B} \left(y _ {i} \log \widehat {y _ {i}} + \left(1 - y _ {i}\right) \log \left(1 - \widehat {y _ {i}}\right)\right). \tag {13}
$$
Table 1: Evaluation performance on GOOD benchmark. - denotes abnormal results caused by under-fitting declared in the leaderboard, and / denotes that the method cannot be applied to this dataset. The best is marked with **boldface** and the second best is with **underline**.
Method	GOOD-HIV ↑	GOOD-ZINC ↓	GOOD-PCBA ↑
scaffold	size	scaffold	size	scaffold	size
covariate	concept	covariate	concept	covariate	concept	covariate	concept	covariate	concept	covariate	concept
ERM	69.55(2.39)	72.48(1.26)	59.19(2.29)	61.91(2.29)	0.1802(0.0174)	0.1301(0.0052)	0.2319(0.0072)	0.1325(0.0085)	17.11(0.34)	21.93(0.74)	17.75(0.46)	15.60(0.55)
IRM	70.17(2.78)	71.78(1.37)	59.94(1.59)	(-)	0.2164(0.0160)	0.1339(0.0043)	0.6984(0.2809)	0.1336(0.0055)	16.89(0.29)	22.37(0.43)	17.68(0.36)	15.82(0.52)
VREx	69.34(3.54)	72.21(1.42)	58.49(2.28)	61.21(2.00)	0.1815(0.0154)	0.1287(0.0053)	0.2270(0.0136)	0.1311(0.0067)	17.10(0.20)	21.65(0.82)	17.80(0.35)	15.85(0.47)
GroupDRO	68.15(2.84)	71.48(1.27)	57.75(2.86)	59.77(1.95)	0.1870(0.0128)	0.1323(0.0041)	0.2377(0.0147)	0.1333(0.0064)	16.55(0.39)	21.91(0.93)	16.74(0.60)	15.21(0.66)
Coral	70.69(2.25)	72.96(1.06)	59.39(2.90)	60.29(2.50)	0.1769(0.0152)	0.1303(0.0057)	0.2292(0.0090)	0.1261(0.0060)	17.00(0.38)	22.00(0.46)	17.83(0.31)	16.88(0.58)
DANN	69.43(2.42)	71.70(0.90)	62.38(2.65)	65.15(3.13)	0.1746(0.0084)	0.1269(0.0042)	0.2326(0.0140)	0.1348(0.0091)	17.20(0.46)	22.03(0.72)	17.71(0.26)	15.78(0.40)
Mixup	70.65(1.86)	71.89(1.73)	59.11(3.11)	62.80(2.43)	0.2066(0.0123)	0.1391(0.0071)	0.2531(0.0150)	0.1547(0.0082)	16.52(0.33)	20.52(0.50)	17.42(0.29)	13.71(0.37)
DIR	68.44(2.51)	71.40(1.48)	57.67(3.75)	74.39(1.45)	0.3682(0.0639)	0.2543(0.0454)	0.4578(0.0412)	0.3146(0.1225)	16.33(0.39)	23.82(0.89)	16.04(1.14)	16.80(1.17)
GSAT	70.07(1.76)	72.51(0.97)	60.73(2.39)	56.96(1.76)	0.1418(0.0077)	0.1066(0.0046)	0.2101(0.0095)	0.1038(0.0030)	16.45(0.17)	20.18(0.74)	17.57(0.40)	13.52(0.90)
GREA	71.98(2.87)	70.76(1.16)	60.11(1.07)	60.96(1.55)	0.1691(0.0159)	0.1157(0.0084)	0.2100(0.0081)	0.1273(0.0044)	16.28(0.34)	20.23(1.53)	17.12(1.01)	13.82(1.18)
CAL	69.12(1.10)	72.49(1.05)	59.34(2.14)	56.16(4.73)	/	/	/	/	15.87(0.69)	18.62(1.21)	16.92(0.72)	13.01(0.65)
DisC	58.85(7.26)	64.82(6.78)	49.33(3.84)	74.11(6.69)	/	/	/	/	/	/	/	/
MoleOOD	69.39(3.43)	69.08(1.35)	58.63(1.78)	55.90(4.93)	0.2752(0.0288)	0.1996(0.0136)	0.3468(0.0366)	0.2275(0.2183)	12.90(0.86)	12.92(1.23)	12.64(0.55)	10.30(0.36)
CIGA	69.40(1.97)	71.65(1.33)	61.81(1.68)	73.62(1.33)	/	/	/	/	/	/	/	/
iMoLD	72.93(2.29)	74.32(1.63)	62.86(2.58)	77.43(1.32)	0.1410(0.0054)	0.1014(0.0040)	0.1863(0.0139)	0.1029(0.0067)	17.32(0.31)	22.58(0.67)	18.02(0.73)	18.21(1.10)
Scoring GNN Regularization. For the scoring  $\mathrm{GNN}_S$ , we apply a penalty to the separation of invariant features, preventing abundance or scarcity and ensuring a reasonable selection. To achieve this, we introduce a regularization term  $\mathcal{L}_{\mathrm{reg}}$  to constrain the size of the selected invariant features:
$$
\mathcal {L} _ {\text {r e g}} = \left| \frac {\langle \mathbf {J} , \mathbf {S} \rangle_ {F}}{| \mathcal {V} | \times d} - \gamma \right|, \tag {14}
$$
where  $\mathbf{J} \in \mathbb{R}^{|\mathcal{V}| \times d}$  denotes a matrix with all entries as  $1, \langle \cdot, \cdot \rangle_F$  is the Frobenius dot product and  $\gamma$  is a pre-defined threshold between 0 and 1.
Codebook Regularization. In the RVQ module, following VQ-VAE [52, 53], to encourage  $\mathbf{h}_v$  to remain close to the selected codebook embedding  $\mathbf{e}_k$  in Equation (5) and prevent it from frequently fluctuating between different embeddings, we add a commitment loss  $\mathcal{L}_{\mathrm{cmt}}$  to foster that  $\mathbf{h}_v$  commits to  $\mathbf{e}_k$  and does not grow uncontrollably:
$$
\mathcal {L} _ {\mathrm {c m t}} = \sum_ {v \in \mathcal {V}} \| \operatorname {s g} [ \mathbf {e} _ {k} ] - \mathbf {h} _ {v} \| _ {2} ^ {2}. \tag {15}
$$
Finally, the learning objective can be defined as the weighted sum of the above losses:
$$
\mathcal {L} = \mathcal {L} _ {\text {p r e d}} + \lambda_ {1} \mathcal {L} _ {\text {i n v}} + \lambda_ {2} \mathcal {L} _ {\text {r e g}} + \lambda_ {3} \mathcal {L} _ {\text {c m t}}, \tag {16}
$$
where  $\lambda_{1},\lambda_{2}$  and  $\lambda_{3}$  are hyper-parameters to control the weights of  $\mathcal{L}_{\mathrm{inv}}$  (in Equation (12)),  $\mathcal{L}_{\mathrm{reg}}$  (in Equation (14)) and  $\mathcal{L}_{\mathrm{cmt}}$  (in Equation (15)), respectively.
In this section, we conduct extensive experiments to answer the research questions. (RQ1) Can our method iMoLD achieve better OOD generalization performance against SOTA baselines? (RQ2) How does each of the component we propose contribute to the final performance? (RQ3) How can we understand the latent discrete space induced by the proposed RVQ?
Datasets. We employ two real-world benchmarks for OOD molecular representation learning. Details of datasets are in Appendix A.
- GOOD [63], which is a systematic benchmark tailored specifically for graph OOD problems. We utilize three molecular datasets for the graph prediction task: (1) GOOD-HIV [16], where the objective is binary classification to predict whether a molecule can inhibit HIV; (2) GOOD-ZINC [64], which is a regression dataset aimed at predicting molecular solubility; and (3) GOOD-PCBA [16], which includes 128 bioassays and forms 128 binary classification tasks. Each dataset comprises two environment-splitting strategies (scaffold and size), and two shift types (covariate and concept) are applied per splitting outcome, resulting in a total of 12 distinct datasets.
Table 2: Evaluation performance on DrugOOD benchmark. The best is marked with boldface and the second best is with underline.
Method	IC50↑	EC50↑
Assay	Scaffold	Size	Assay	Scaffold	Size
ERM	71.63(0.76)	68.79(0.47)	67.50(0.38)	67.39(2.90)	64.98(1.29)	65.10(0.38)
IRM	71.15(0.57)	67.22(0.62)	61.58(0.58)	67.77(2.71)	63.86(1.36)	59.19(0.83)
Coral	71.28(0.91)	68.36(0.61)	64.53(0.32)	72.08(2.80)	64.83(1.64)	58.47(0.43)
MixUp	71.49(1.08)	68.59(0.27)	67.79(0.39)	67.81(4.06)	65.77(1.83)	65.77(0.60)
DIR	69.84(1.41)	66.33(0.65)	62.92(1.89)	65.81(2.93)	63.76(3.22)	61.56(4.23)
GSAT	70.59(0.43)	66.45(0.50)	66.70(0.37)	73.82(2.62)	64.25(0.63)	62.65(1.79)
GREA	70.23(1.17)	67.02(0.28)	66.59(0.56)	74.17(1.47)	64.50(0.78)	62.81(1.54)
CAL	70.09(1.03)	65.90(1.04)	66.42(0.50)	74.54(4.18)	65.19(0.87)	61.21(1.76)
DisC	61.40(2.56)	62.70(2.11)	61.43(1.06)	63.71(5.56)	60.57(2.27)	57.38(2.48)
MoleOOD	71.62(0.52)	68.58(1.14)	65.62(0.77)	72.69(1.46)	65.74(1.47)	65.51(1.24)
CIGA	71.86(1.37)	69.14(0.70)	66.92(0.54)	69.15(5.79)	67.32(1.35)	65.65(0.82)
iMoLD	72.11(0.51)	68.84(0.58)	67.92(0.43)	77.48(1.70)	67.79(0.88)	67.09(0.91)
- DrugOOD [13], which is a OOD benchmark for AI-aided drug discovery. This benchmark provides three environment-splitting strategies, including assay, scaffold and size, and applies these three splitting to two measurements (IC50 and EC50). As a result, we obtain 6 datasets, and each dataset contains a binary classification task for drug target binding affinity prediction. A detailed description of each environment-splitting strategy is also in Appendix A.
Baselines. We thoroughly compare our method against ERM [65] and two groups of OOD baselines: (1) general OOD algorithms used for Euclidean data, which include domain adaptation methods such as Coral [44] and DANN [43], group distributionally robust optimization method (Group-DRO [41]), invariant learning methods such as IRM [19] and VREx [40] and data augmentation method (Mixup [66]). And (2) graph-specific algorithms, including Graph OOD algorithms such as CAL [28], DisC [32], MoleOOD [25] and CIGA [26], as well as interpretable graph learning methods such as DIR [29], GSAT [49] and GREA [47]. Details of baselines and implementation are in Appendix B.
Evaluation. We report the ROC-AUC score for GOOD-HIV and DrugOOD datasets as the task is binary classification. For GOOD-ZINC, we use the Mean Average Error (MAE) since the task is regression. While for GOOD-PCBA, we use Average Precision (AP) averaged over all tasks as the evaluation metric due to extremely imbalanced classes. We run experiments 10 times with different random seeds, select models based on the validation performance and report the mean and standard deviations on the test set.
Table 1 and Table 2 present the empirical results on GOOD and DrugOOD benchmarks, respectively. Our method iMoLD achieves the best performance on 16 of the 18 datasets and ranks second on the other two datasets. Among the compared baselines, the graph-specific OOD methods perform best on only 11 datasets, and some general OOD methods outperform them on another 7 datasets. This suggests that although some advanced graph-specific OOD methods can achieve superior performance on some synthetic datasets (e.g., to predict whether a specific motif is present in a synthetic graph [29]), they may not perform well on molecules due to the realistic and complex data structures and distribution shifts. In contrast, our method is able to achieve the best performance on most of the datasets, which indicates that the proposed identification of invariant features in the latent space is effective for applying the invariance principle to the molecular structure. We also observe that MoleOOD, a method designed specifically for molecules, does not perform well on GOOD-ZINC and GOOD-PCBA, possibly due to its dependence on inferred environments. Inferring environments may become more challenging for larger-scale datasets, such as GOOD-ZINC and GOOD-PCBA, which contain hundreds of thousands of data and more complex tasks (e.g., PCBA has a total of 128 classification tasks). Our method does not require the inference of environment, and is shown to be effective on datasets of diverse scales and tasks.
(a) ERM
Score(train) = 90.08, Score(val) = 78.57  
D(Y=0) = 0.46, D(Y=1) = 2.41
Figure 4: Visualization of the extracted features on training and validation set when the model achieves the highest score on the validation set. Score(train) and Score(val) are ROC-AUC scores on the training and validation set, respectively.  $D(Y = 0)$  and  $D(Y = 1)$  are distances between the features on the training and validation sets of each class.
(b)ERM  $(+VQ)$
Score(train) = 86.46, Score(val) = 79.03  
D(Y=0) = 0.28, D(Y=1) = 0.91
(c) ERM (+RVQ)
Score(train) = 92.70, Score(val) = 78.96  $D(Y = 0) = 0.20$ $D(Y = 1) = 1.68$
(d) iMoLD
Score(train) = 97.69, Score(val) = 79.89  
D(Y=0) = 0.19, D(Y=1) = 0.64
5.3 Ablation Studies of Components (RQ2)
We conduct ablation studies to analyze how each proposed component contributes to the final performance. Three groups of variants are included: (1) The first group is to investigate the RVQ module. It consists of ERM and ERM (+RVQ), to explore the performance of the simplest ERM and the role of applying the RVQ to it. It also consists of variants of our method, including w/o VQ, which uses H instead of  $\mathbf{H}^{\prime}$  in Equation (9) to obtain  $\mathbf{H}^{\mathrm{inv}}$  and  $\mathbf{H}^{\mathrm{Spu}}$ ; w/o R, which obtains  $\mathbf{H}^{\mathrm{Inv}}$  and  $\mathbf{H}^{\mathrm{Spu}}$  without a residual connection by using the output of discretization in Equation (5). (2) The second group is to investigate the impact of each individual training objective, including w/o  $\mathcal{L}_{\mathrm{inv}}$ , w/o  $\mathcal{L}_{\mathrm{reg}}$  and w/o  $\mathcal{L}_{\mathrm{cnt}}$ , which remove  $\mathcal{L}_{\mathrm{inv}}$ ,  $\mathcal{L}_{\mathrm{reg}}$  and  $\mathcal{L}_{\mathrm{cnt}}$  from the final objective function in Equation (16), respectively. (3) The third group consists of variants that replace our proposed self-supervised invariant learning objective with other methods, including w/  $\mathcal{L}_{CIGA}$  and w/  $\mathcal{L}_{GREA}$ , which represent the objective functions using CIGA [26] and GREA [47], respectively.
Table 3 reports the results on the scaffold-split of GOOD-HIV and the size-split of GOOD-PCBA. We observe that the addition of the RVQ module can improve performance without using any invariant learning algorithms. This indicates that the designed RVQ is effective. Notably, on the GOOD-HIV dataset, removing the VQ module results in significant performance degradation. While on the GOOD-PCBA dataset, removing the residual connection leads to a significant performance drop. This observation may be attributed to the larger scale and more complex task of the GOOD-PCBA dataset. Although VQ im
Table 3: Ablation studies. / denotes that the variant cannot be applied to this dataset.
GOOD-HIV-Scaffold	GOOD-PCBA-Size
covariate	concept	covariate	concept
iMoLD	72.93(2.29)	74.32(1.63)	18.02(0.73)	18.21(1.10)
ERM	69.55(2.39)	72.48(1.26)	17.75(0.46)	15.60(0.55)
ERM (+RVQ)	70.18(1.07)	73.14(1.99)	17.84(0.47)	16.59(0.40)
w/o VQ	70.18(1.60)	72.88(0.66)	17.88(1.57)	17.37(0.82)
w/o R	72.40(2.21)	73.03(0.15)	17.37(0.62)	12.99(2.09)
w/o Linv	70.73(1.35)	71.89(1.80)	17.55(0.34)	17.05(0.87)
w/o Lleg	71.47(0.72)	72.39(0.87)	17.33(0.68)	17.91(0.34)
w/o Lcnt	70.20(4.07)	71.98(1.36)	17.24(1.02)	17.68(2.20)
w/ LCIIGA	71.23(1.31)	73.31(1.73)	/	/
w/ LGREA	72.33(1.23)	72.94(2.30)	17.93(0.55)	16.86(0.50)
proves the generalization, it reduces the expressivity capacity of the model. Thus, achieving a balance between the benefits and weaknesses brought by VQ is necessary with the proposed RVQ module. By observing the second group, we can conclude that each module is instrumental to the final performance. Removing the invariant learning objective  $\mathcal{L}_{\mathrm{inv}}$  and commitment loss  $\mathcal{L}_{\mathrm{cnt}}$  results in more pronounced performance degradation, suggesting that the invariant learning objective is effective and indispensable, as well as the need for a constraint term to justify the output of the encoder and the codebook in VQ. By comparing the third group of results, we find that our proposed learning objective outperforms the existing methods, which shows that our approach can not only be applied to a wider range of tasks but also achieve better performance. Moreover, to investigate the robustness of our model, we perform a hyper-parameter sensitivity analysis in Appendix C.1.
5.4 Analysis of Latent Discrete Space (RQ3)
To further explore why our method performs better and to understand the superiority of the introduced discrete latent space, we visualize the projection of extracted  $\mathbf{z}^{\mathrm{Inv}}$  on both training and validation set
when the model achieves the best score on the validation set, using t-SNE [67] on the covariate-shift dataset of GOOD-HIV-Scaffold in Figure 4 (d). We also visualize the results of some baselines, including the vanilla ERM (Figure 4 (a)), the ERM equipped with VQ after encoder (ERM(+VQ), Figure 4 (b)), and with the RVQ module (ERM(+RVQ), Figure 4 (c)). We additionally compute the 1-order Wasserstein distance [68] between the features on the training and validation sets of each class, to quantify the dissimilarity in the feature distribution across varying environments. We find that adding the VQ or RVQ after the encoder results in a more uniform distribution of features and lower feature distances, due to the fact that VQ makes it possible to reuse previously encountered embeddings in new environments by discretizing them. Moreover, the feature distribution of our method is more uniform and the distance is smaller. This suggests that our method is effective in identifying features that are invariant across different environments. Moreover, we also observe that all ERM methods achieve lower validation scores with higher training scores, implying that these methods are prone to overfitting. Our method, on the other hand, achieves not only a higher validation score but also a higher corresponding training score, thereby demonstrating its ability to overcome the problem of easy overfitting and improve the generalization ability effectively.
In this work, we propose a new framework that learns invariant molecular representation against distribution shifts. We adopt a "first-encoding-then-separation" strategy, wherein a combination of encoding GNN and residual vector quantization is utilized to derive molecular representation in latent discrete space. Then we learn a scoring GNN to identify invariant features from this representation. Moreover, we design a task-agnostic self-supervised learning objective to enable precise invariance identification and versatile applicability to various tasks. Extensive experiments on real-world datasets demonstrate the superiority of our method on the molecular OOD problem. Overall, our proposed framework presents a promising approach for learning invariant molecular representations and offers valuable insights for addressing distribution shifts in molecular data analysis.
This work was supported by the National Key Research and Development Program of China (2022YFB4500300), National Natural Science Foundation of China (NSFCU19B2027, NSFC91846204, NSFC62302433), joint project DH-2022ZY0012 from Donghai Lab, and sponsored by CCF-Tencent Open Fund (CCF-Tencent RAGR20230122). We want to express gratitude to the anonymous reviewers for their hard work and kind comments.
[1] Eugene N Muratov, Jürgen Bajorath, Robert P Sheridan, Igor V Tetko, Dmitry Filimonov, Vladimir Poroikov, Tudor I Oprea, Igor I Baskin, Alexandre Varnek, Adrian Roitberg, et al. Qsar without borders. Chemical Society Reviews, 49(11):3525-3564, 2020.
[2] Stephani Joy Y Macalino, Vijayakumar Gosu, Sunhye Hong, and Sun Choi. Role of computer-aided drug design in modern drug discovery. Archives of pharmacal research, 38:1686-1701, 2015.
[3] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In ICML, volume 70 of Proceedings of Machine Learning Research, pages 1263-1272. PMLR, 2017.
[4] Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman-Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8):3370-3388, 2019.
[5] Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. In NeurIPS, 2020.
[56] Jun Xia, Chengshuai Zhao, Bozhen Hu, Zhangyang Gao, Cheng Tan, Yue Liu, Siyuan Li, and Stan Z. Li. Mole-BERT: Rethinking pre-training graph neural networks for molecules. In The Eleventh International Conference on Learning Representations, 2023.
[57] Dianbo Liu, Alex Lamb, Kenji Kawaguchi, Anirudh Goyal, Chen Sun, Michael C. Mozer, and Yoshua Bengio. Discrete-valued neural communication. In NeurIPS, pages 2109-2121, 2021.
[58] Frederik Träuble, Anirudh Goyal, Nasim Rahaman, Michael Mozer, Kenji Kawaguchi, Yoshua Bengio, and Bernhard Schölkopf. Discrete key-value bottleneck. CoRR, abs/2207.11240, 2022.
[59] Dianbo Liu, Alex Lamb, Xu Ji, Pascal Notsawo, Michael Mozer, Yoshua Bengio, and Kenji Kawaguchi. Adaptive discrete communication bottlenecks with dynamic vector quantization. CoRR, abs/2202.01334, 2022.
[60] Jose G Moreno-Torres, Troy Raeder, Rocio Alaiz-Rodriguez, Nitesh V Chawla, and Francisco Herrera. A unifying view on dataset shift in classification. Pattern recognition, 45(1):521-530, 2012.
[61] Joaquin Quinonero-Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008.
[62] Gerhard Widmer and Miroslav Kubat. Learning in the presence of concept drift and hidden contexts. Machine learning, 23:69-101, 1996.
[63] Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. GOOD: A graph out-of-distribution benchmark. CoRR, abs/2206.08452, 2022.
[64] Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuu-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science, 4(2):268-276, 2018.
[65] Vladimir Cherkassky. The nature of statistical learning theory. IEEE Trans. Neural Networks, 8(6):1564, 1997.
[66] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In ICLR (Poster). OpenReview.net, 2018.
[67] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
[68] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
[69] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In NeurIPS, pages 8024-8035, 2019.
[70] Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. CoRR, abs/1903.02428, 2019.
In this paper, we use 18 publicly benchmark datasets, 12 of which are from GOOD [63] benchmark. They are the combination of 3 datasets (GOOD-HIV, GOOD-ZINC and GOOD-PCBA), 2 types of distribution shift (covariate and concept), and 2 environment-splitting strategies (scaffold and size). The rest 6 datasets are from DrugOOD [13] benchmark, including IC50-Assay, IC50-Scaffold, IC50-Size, EC50-Assay, EC50-Scaffold, and EC50-Size. The prefix denotes the measurement and the suffix denotes the environment-splitting strategies. This benchmark exclusively focuses on covariate shift. We use the latest data released on the official webpage<sup>3</sup> based on the ChEMBL 30 database<sup>4</sup>. We use the default dataset split proposed in each benchmark. For covariate shift, the training, validation and testing sets are obtained based on environments without interactions. For concept shift, a screening approach is leveraged to scan and select molecules in the dataset. Statistics of each dataset are in Table 4.
Table 4: Dataset statistics.
Dataset	Task	Metric	#Train	#Val	#Test	#Tasks
GOOD	HIV	scaffold	covariate	Binary Classification	ROC-AUC	24682	4133	4108	1
concept	Binary Classification	ROC-AUC	15209	9365	10037	1
size	covariate	Binary Classification	ROC-AUC	26169	4112	3961	1
concept	Binary Classification	ROC-AUC	14454	3096	10525	1
ZINC	scaffold	covariate	Regression	MAE	149674	24945	24946	1
concept	Regression	MAE	101867	43539	60393	1
size	covariate	Regression	MAE	161893	24945	17042	1
concept	Regression	MAE	89418	19161	70306	1
PCBA	scaffold	covariate	Multi-task Binary Classification	AP	262764	44019	43562	128
concept	Multi-task Binary Classification	AP	159158	90740	119821	128
size	covariate	Multi-task Binary Classification	AP	269990	43792	31925	128
concept	Multi-task Binary Classification	AP	150121	32168	115205	128
DrugOOD	IC50	assay	Binary Classification	ROC-AUC	34953	19475	19463	1
scaffold	Binary Classification	ROC-AUC	22025	19478	19480	1
size	Binary Classification	ROC-AUC	37497	17987	16761	1
EC50	assay	Binary Classification	ROC-AUC	4978	2761	2725	1
scaffold	Binary Classification	ROC-AUC	2743	2723	2762	1
size	Binary Classification	ROC-AUC	5189	2495	2505	1
A.2 The Cause of Molecular Distribution Shift
The molecule data can be divided according to different environments, and distribution shifts occur when the source environments of data are different during training and testing. In this work, we investigate three types of environment-splitting strategies, i.e., scaffold, size and assay. And the explanation of each environment are in Table 5.
B.1Baselines
We adopt the following methods as baselines for comparison, one group of which are common approaches for non-Euclidean data: