Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets
Authors: Dominique Beaini
Shenyang Huang
Joao Alex Cunha
Zhiyi Li
Gabriela Moisescu-Pareja
Oleksandr Dymov
Samuel Maddrell-Mander
Callum McLean
Frederik Wenkel
Luis Müller
Jama Hussein Mohamud
Ali Parviz
Michael Craig
Michał Koziarski
Jiarui Lu
Zhaocheng Zhu
Cristian Gabellini
Kerstin Klaser
Josef Dean
Cas Wognum
Maciej Sypetkowski
Guillaume Rabusseau
Reihaneh Rabbany
Jian Tang
Christopher Morris
Mirco Ravanelli
Guy Wolf
Prudencio Tossou
Hadrien Mary
Therence Bois
Andrew W Fitzgibbon
Blazej Banaszewski
Chad Martin
Dominic Masters
Venue: ICLR 2024 poster (2024)
Keywords: graph neural networks
Datasets
molecules
molecular graphs
Quantum
Multi-task
foundation model
Primary area: learning on graphs and other geometries & topologies
TL;DR: By far the largest datasets for 2D molecular graphs, both in terms of number of unique molecules and number of tasks, and a new library for multi-task GNNs.

Abstract: Recently, pre-trained foundation models have enabled significant advancements in multiple fields. In molecular machine learning, however, where datasets are often hand-curated, and hence typically small, the lack of datasets with labeled features, and codebases to manage those datasets, has hindered the development of foundation models. In this work, we present seven novel datasets categorized by size into three distinct categories: ToyMix, LargeMix and UltraLarge. These datasets push the boundaries in both the scale and the diversity of supervised labels for molecular learning. They cover nearly 100 million molecules and over 3000 sparsely defined tasks, totaling more than 13 billion individual labels of both quantum and biological nature. In comparison, our datasets contain 300 times more data points than the widely used OGB-LSC PCQM4Mv2 dataset, and 13 times more than the quantum-only QM1B dataset. In addition, to support the development of foundational models based on our proposed datasets, we present the Graphium graph machine learning library which simplifies the process of building and training molecular machine learning models for multi-task and multi-level molecular datasets. Finally, we present a range of baseline results as a starting point of multi-task and multi-level training on these datasets. Empirically, we observe that performance on low-resource biological datasets show improvement by also training on large amounts of quantum data. This indicates that there may be potential in multi-task and multi-level training of a foundation model and fine-tuning it to resource-constrained downstream tasks. The Graphium library is publicly available on Github and the dataset links are available in Part 1 and Part 2.

Published as a conference paper at ICLR 2024
TOWARDS FOUNDATIONAL MODELS FOR MOLECULAR LEARNING ON LARGE-SCALL MULTI-TASK DATASETS
Dominique Beaini'.
Shenyang Huang Joao Alex Cunha Zhiyi Li
Gabriela Moisescu-Pareja--e
Oleksandr Dymov! Samuel Maddrell-Mander
Callum McLean
Frederik Wenkel Luis Müller? Jama Hussein Mohamud'
Ali Parviz!.20 12, Michael Craig² Michal Koziarski Jiarui Lu!.3 3
Zhaocheng Zhu Cristian Gabellini Kerstin Klaser Josef Dean
Cas W Maciej Sypetkowski
Guillaume Rabusseaul.3,
Reihaneh Rabbany 1.99 Jian T
Christopher Morris' Mirco Ravanelli
Guy Wolf Prudencio Tossou Hadrien Mary²
Therence Bois"
Andrew Fitzgibbon Blažej Banaszewski
Chad Martin
Dominic Masters
Mila- Québec AI Institute
Valence Labs
Université de Montréal
McGiil University
'Graphcore RWTH Aachen University
NNw Jersey Institute of Technology
SHEC Montréal
'CIFAR AI Chair
ABSTRACT
Recently, pre-trained foundation models have enabled significant advancements in multiple fields. In molecular machine learning, however, where datasets are often hand-curated, and hence typically small, the lack of datasets with labeled features, and codebases to manage those datasets, has hindered the development of foundation models. In this work, present seven novel datasets categorized by size into three distinct categories: ToyMix, LargeMix and UltraLarge. These datasets push the boundaries in both the scale and the diversity of supervised labels for molecular learning. They cover nearly 100 million molecules and over 3000 sparsely defined tasks, totaling more than 13 billion individual labels of both quantum and biological nature. In comparison, our datasets contain 300 times more data points than the widely used OGB-LSC PCQM4Mv2 dataset, and 13 times more than the quantum-on QMIB dataset. In addition, to support the development of foundational models based on our proposed datasets, we present the Graphium graph machine learning library which simplifies the process of building and training molecular machine learning models for multi-task and multi-level molecular datasets. Finally, we present range of baseline results as starting point of multi-task and multi-level training on these datasets. Empirically, we observe that performance on low-resource biological datasets shows improvement by also training on large amounts of quantum data. This indicates that there may be potential in multi-task and multi-level training of a foundation model and fine-tuning it to resource-constrainee downstream tasks. The Graphium library is publicly available on Github and the dataset links available in Part and Part 2
INTRODUCTION
Graph and geometric deep learning models have been key component in the recent successes of machine learning in drug discovery (Gasteiger al., Masters et al., 2023 Rampášek et al. 2022 Ying et al. 2021). These methods have demonstrated considerable performance in molecular representation learning (in the 2D (Rampásek et al. 2022), 3D (Gasteiger et al.; 2021) ) and 4D (Wu et al. 2022 cases), activity and property prediction (Huang et al. 2021), I), force field development (Batatia et al. 2022), molecular generation (Bilodeau et al., 2022). 2, and the modeling of atomistic interactions (Corso al.) (022). Like other deep learning methods, they require significant training data for high modeling accuracy. However, in the current therapeutics literature, most training datasets have limited samples (Huang al. 202T). Tantalizingly, recent progress in self-supervised learning and

---

Published as a conference paper at ICLR 2024
foundation models in natural language processing (NLP) (Brown et al. 2020a Liu et al. 2023b and computer vision (CV) (Dosovitskiy et al.) has drastically improved data efficiency in deep learning. In fact. by investing upfront in pre-training large models with lots of data, one-time cost, it is proven that the learned inductive bias lowers the data requirements for downstream tasks.
Following these successes, many studies have explored the pre-training of large molecular graph neural networks and their benefits in low-data molecular modeling (Lin et al. 2022 Méndez-Lucio
al. 2022, Zhou et al. 2023). However, due the scarcity of large and labeled molecular datasets, these studies could only leverage self-supervised techniques such as contrastive learning, auto- encoders, or denoising tasks (Hu et al.) 2020a] Xia et al. 2023). Low-data modeling efforts by fine-tuning from these models have, as yet. yielded just fraction of the advancement achieved by self-supervised models NLP and CV (Sun et al. 2022). This is partially explained by the underspecification of molecules, and their conformers, as graphs, as their behavior is environment- dependent and mainly governed by quantum physics. For example, it is well known that structurally similar molecules may have vastly different bioactivity. known activity cliffs, and this limits graph modeling based only on structural information (van Tilborg et al. 2022). We argue that building effective foundation models for molecular modeling requires supervised training with both quantum mechanical (QM) descriptions and biological environment-dependee data.
1.1 CONTRIBUTIONS
This study advances molecular research in three ways. First, we introduce a new family of multi-task datasets, orders of magnitude larger than the current state-of-the-art. Second, we describe graph machine learning library called Craphi um to facilitate efficient training on these extensive datasets. Third, range of baseline models are implemented, supporting the value of training on diverse collection of tasks.
Datasets We introduce three extensive and meticulously curated multi-label datasets that cover nearly 100 million molecules and over 3000 sparsely defined tasks, totaling more than 13 billion individual labels, currently the largest of their kind. These datasets have been designed for the supervised training of foundation models by combining labels representing quantum and biological properties acquired through both simulation and wet lab experimentation. The labels are also multi-level, encompassing both node-level and graph-level tasks. The diversity of labels facilitates efficient transfer learning and enables the construction of foundational models by improving their generalization ability for a wide range of downstream molecular modeling tasks.
To create these comprehensive datasets, we carefully curated and enhanced existing data with additional information. As result, each molecule in our collection is accompanied by descriptions of its quantum mechanical (QM) properties and/or biological activities. The QM properties cover energetic, electronic, and geometric aspects, computed using various advanced methods, including density functional theory (DFT) methods like B3LYP (Nakata & Shimazaki, 2017). as well as semi-empirical methods like PM6 (Nakata et al. 2220). On the biological activity side, datasets include molecular signatures obtained from dose-response bioassays, gene expression profiling, and toxicological profiling, depicted in Figure [ The joint modeling of quantum and biological effects promotes the ability to describe complex environment-depentee be infeasible to extract from what are typically limited experimental datasets.
properties of molecules that would
The Graphium Library We have developed a comprehensive graph machine learning library called Graphi um to facilitate efficient training on these extensive multi-task datasets. This novel library simplifies the process of building and training molecular graph foundation models by incorporating feature ensembles and complex feature interactions. By treating features (both positional and structural) and representations as fundamental building blocks and implementing state-of-the-ar GNN layers, Graphi i um overcomes the challenges posed by existing frameworks that were primarily designed for sequential samples with limited interactions among node, edge, and graph features. Additionally, by providing features such as dataset combination, missing data handling, and joint training, Gr Graphin um handles the critical and otherwise complex engineering of training models on large dataset ensembles in a straightforward and highly customizable way.
Baseline Results We train range of models in both single-dataset and multi-datase scenarios for the dataset mixes presented. These provide solid baselines that can help guide future users of these datasets, but also give some indication of the value of training in this multi-dataset paradigm.
2

---

Published as a conference paper at ICLR 2024
PCOMAM_E22 CZ5, N4
PCRA_1928
11000, VCAP
Fooetigg
Minenn cats
1.501 mik.
Nevoral class
STR GT8---ooek 11000 MO7
G-labels
Giatel
TorMix
1,308 G-libel:
315
N: N-labes
10
Made- bxtels
LARGEMIX
$78
62 alet, 7 Rubl ULTRALARGE
Visual summary of the proposed collections of molecular datasets. The "mixes" are meant be predicted simultaneously multi-task fashion. They include quantum, chemical, and biological properties, categorical and continuous data points, graph-level and node-level tasks.
Figure 1:
Specifically, the results on these models show that training of low-resource tasks can be significantly improved by training combination with larger datasets.
In summary, our study introduces the largest 2D molecular datasets to date, specifically designed to train foundation models that can effectively comprehend molecules' quantum properties and biological adaptability and therefore be fine-tuned to wide range of downstream tasks. In addition, we have developed the Graphi um library streamline the training process these models and show a range of baseline results that highlight the effectiveness of both the datasets and library presented.
1.2 RELATED WORK
Molecular Datasets Several molecular datasets are available for node- and graph-level supervised learning tasks. A few notable examples are MoleculeNet (Wu al.) 2018). which is a collection of smaller datasets; QM9, which consists of quantum properties of 134k molecular conformers (Ramakrishnan et al. 2014, Wu et al.) 2018); ZINC comprising the difference between the partition coefficient (clogP) P) and the synthetic accessibility score (SAS) of 12k molecules (Dwivedi et al.) 2020); PCQM4M(V2)) including the HOMO-LUMO gap (the lowest excitation energy of molecule) of 3.6M graphs (Hu et al. 2021b). Furthermore, Open Graph Benchmark's s (OGB) (Hu et al. 2020b) large molecular datasets are OGBG- (~40k) and (GBG-MOLPCHE (~>437k). And PCQM4M, the largest dataset, is still inadequate for large-scale supervised molecular pre-traie pre-training, especially since only contains a single property per molecule.
Several examples of large-scale datasets for self-supervised training exist, notably the 1.2 million molecule MOLACAMO.BBrrr et al. 2019) dataset, and the ~ 880 million molecule ZINC20 dataset (Iwwin 2020), and while these provide a valuable training platform they lack experimental labels. Examples of curated self-supervised datasets of 72 million molecules are used in Winter et al. (2019), however these datasets typically computable quantum chemical properties rather than enviromment-depenn dependent labels. Some work in this direction, as proposed by Mayr et al. (2018). shows supervised datasets can be extracted from the large CHEMBL dataset, and in Mundez-Lucio et al.) (2022) dataset of some 0.5 million molecules with over 1000 labels is produced. Recent work on accelerating DFT calculations in Mathiasen et al.) (2023) offers a tantalizing avenue to increasing the size of available 3D quantum datasets with billion data points, albeit quantum-only, Finally, the GDB 17 dataset (Ruddigkeit et al. (0122) does enumerate some 166 billion chemically possible molecules, but does not provide labels or features associated with them.
However, none of these datasets are both sufficiently large, and provided with enough high quality experimental labels, to reflect the indirect properties relevant to developing foundation model for drug discovery.
As such, dataset collection departs from previous ones it encompasses multiple graph-level and node-level tasks, in addition to complex anvironment-dependeen properties.
Pre-trained Chemical Foundation Models Recent works explored pre-training for molecular property prediction via self-supervised learning techniques such as denoising (Zaidi et al., 2022); auto-encoders (Honda et al. 2019) and contrastive learning (Suresh et al. 2021). For comprehensive
3

---

Published as a conference paper at ICLR 2024
overview of supe-vsd approaches pre-training on molecular graphs, see la et al. 2023). As already discussed in Section [ self-supervised pre-training does not take into account the quantum and biochemical characteristies of the molecules, necessitating pre- -re-traiiin in supervised fashion. There have been some efforts at few-shot learning (Stanley et al. 2021 Schimunek et al. 2023). but the applicability of such methods is still limited both by a proper molecular featurization and a proper task featurization, the former of which we aim to tackle in our current work.
Graph Learning Libraries In recent years, many libraries have emerged to speed up the development of graph-learning models. PyTTrrh Geometric (PyG) (ee & Lenssen 2019) is built on Py Torch and provides essential functions for graph machine learning including a custom data format for fast processing and loading of graph data, as well as various model implementations. DGL (Wang et al. 2019 is general- purposse graph learning library that works with PyTorch, TensorFlow, or Jax and supports various model implementations and custom data loading format. Jraph (Godwin R et al. 2020) for Jax, tf ff_eemettric (Hu et al. a for TensorFlow, StellarGraph (Data61, 2018) built on Keras, GraphNets (Battaglia et al., 2018) for TensorFlow and Sonnet, and CogDL (Cen et al. 2023) are other libraries that provide building blocks for graph learning operations. In addition, GraphGym (You et al., (220)) is an experimental platform for graph learning experiments and model design, based on PyG. Although these libraries provide the fundamental components for creating and training graph learning models, extra engineering work is essential to construct scalable training and inference pipelines. In addition, there are specialized libraries that cater to drug discovery, including TorchDrug (Zhu et al.) 2022) and DeepChem (Ramsundar al., 2019). At the same time, GraphGPS (Rampášek et al. 2022) is dedicated to experiment and model design for graph transformers, and DIG (Liu et al., (022) offers toolkit for experimental design and model evaluation.
However, training and inference across multiple tasks and levels are yet to be explored in the context of molecular learning. Furthermore, existing libraries fail tackle the difficulties of pre-training large-scale models, such as loading and managing large amounts of graph data efficiently. Thus, our datasets are accompanied by the Graphi um library to facilitate the use of state-of-the- GNNs in a large-scale multi-task setting. Graphi um enables researchers to effortlessly apply GNN models to large-scale molecular datasets while paving the way towards foundational GNNs for chemistry.
2 PROPOSED DATASETS
Instead of proposing a single task dataset, we suggest a collection or "mix" of datasets enabling multi- task learning across relevant chemical tasks helpful in learning quantum chemistry and biochemistry and building foundation models. The datasets are visualised in Figure [ complemented with statistics in Table See Appendix B.2 and the Reproducibility section for details on datasets' licenses and availability, and see Appendix D for deeper dive into the datasets.
Firstly, we propose smaller TOYMIX dataset to enable fast iterations of the architecture while also providing valuable insight. Secondly, we propose the LARGEMIX, mix of well-curated datasets with data for millions of chemical compounds over thousands of biological or quantum labels. Finally, we propose the ultra-large PM6 83M dataset, which pushes the limit of dataset scale with over 12 billion labels defined for 83 million molecules. In the following section, the three dataset categories described in detail.
Table 1: Dataset statistics. /. denote graph-level / node-level tasks, respectively, and C., R., and RC. denote classification, regression and ranked classification respectively. See Table [orr detailed statistics.
duta points
N. data
Dataset
Type
Task
mol.
G. labels
N. labels
points
QM9
Quantum
R.
134k
19
2.5M
Mix
ZiNC12K
Computed
Tox21
Bio-assays
C.
8k
12
75k
PCOM4M_G23
Quantum
R.
3.8M
25
93M
197.7M
LARGEMIx
PCBA_132
Bio-assays
1.6M
1,328
224.4M
L1000_VCAR
Tramscriptomios
RC.
15%
978
15M
1000_MCFF
Transcriptomics
RC.
12k
978
11M
UUTRALARGE
PM6 83M
Quantum
R.
83M
62
4,003M
7
8,509.2M
4

---

Published as a conference paper at ICLR 2024
2.1 ToYMIx (QM9. Tox21 + ZINC12K)
The ToYMIX dataset combines the QM9, Tox21, and ZINC12K datasets. These datasets are well-known in the literature and used as toy datasets, or very simple datasets, in various contexts to make it easy to iterate quickly on models. By regrouping toy datasets from quantum ML, drug discovery, and GNN expressivity, we hope that the learned model will be representative of the model performance we can expect the larger datasets. As can be seen in Figure the ToyMIX collection contains 34 labels per molecule across the three datasets, containing a mix of floating point properties and class labels.
Train/Valhe Seaiin/aaaaaatonTeeeee Splits for all the datasets in TOYMIX are split randomly with a ratio of 0.8/0.1/0.1. Random splitting, rather than the more complex schemes we discuss below, is used since it is the simplest and fits the idea of having a toy dataset well.
2021 Gasteiger et al.). It QM9 well-known dataset the field of 3D GNNs (Gasteiger al.
consists of 19 graph-level quantum properties associated to an energy-minimized 3D conformation of the molecules (Ramakrishnan et al.) 2014). It is considered simple dataset since all the molecules have at most 9 heavy atoms. We chose QM9 in our TOYMIX since it is very similar to the larger proposed quantum datasets, ICCMtaMmmlltitask and PM6 636M, but with smaller molecules.
multitask
Tox21 is a well-known dataset for researchers in machine learning for drug discovery (Huang et al. 2021, Tice et al., 2013). consists multi-label classification task with 12 labels, with most labels missing and a strong imbalance towards the negative class. We chose Tox21 I in our TOYMIX since it very similar to the larger proposed bioassay dataset, PCBA_ 132__132____ 1728_15645 1564K both in terms of sparsity and imbalance and to the L1000 datasets in terms of imbalance.
ZINC12k a commonly used dataset in the GNN expressivity literature, especially for evaluating positional encodings. graph Transformers. and message passing for complexes (Beaini al. 2021 Dwivedi et al., 2020; Rampásek et al.) 2022] Bodnar et al., 202T). We include it in our ToYMIX as substructure understanding and positional encodings important for generalization, especially in the case of graph Transformers where structural biases are missing. Hence, we hope that performance on this task will correlate well with performance when scaling graph Transformers.
2.2 LARGEMIX (PCQM4M + PCBA1328 + L1000)
In this section, we present the LARGEMIX dataset, comprised of four different datasets with tasks taken from quantum chemistry (PCQM4M), bio-assays (PCBA) and transcriptomics.
Train/validatiion/essttttt seen Splits For the PCQM4M_G22 _4, we create 0.92/0.04/0.Cm 04 split. Then, for all the other datasets in LARGEMIX, we first create a "test_seen" split by taking the set of molecules from L1000 and PCBA1328 that also present in the training set of PCQM4M G25 N4, such that we can evaluate whether having the quantum properties of molecule helps generalize for biological properties. For the remaining parts, split randomly with ratio of 0.92/0.04/0.044.
L1000 VCAP and MCF7 The LINCS L1000 is database generated using high- throughput transcriptomics. More than 30,000 perturbations on set of 978 landmark genes (Subramanian et al., 2017) from multiple cell lines were screened. VCAP and MCF7 are, respectively, prostate cancer and human breast cancer cell lines. In L1000, most of the perturbagens are chemical, meaning that small drug-like molecules are added to the cell lines to observe how the gene expressions change. This allows the generation of biological signatures of the molecules, which are known to correlate with drug activity and side effects (Wang et al., 2016). To process the data into our two datasets comprising the VCAP and MCF7 cell lines, we used their "level 5" data composed of the cleanup data converted to z-scores (Subramanian et al., 2017). and applied a filter to keep only chemical perturbagens. However, we were left with multiple data points per molecule since some variables could change (e.g., incubation time) and generate a new measure. Given our objective of generating a single signature per molecule, we decided to take the measurements with the strongest global activity such that the variance over the 978 genes is maximal. Then, since these signatures are generally noisy, binned them into five classes corresponding to z-scores based the thresholds +2.
5

---

Published as a conference paper at ICLR 2024
The cell lines VCAP and MCF7 were selected since they have higher number of unique molecule perturbagens than other cell lines. They also have a relatively low data imbalance, with ~92% falling in the "neutral class" (where the z-score is between -2 and 2).
PCBA1328 This dataset is very similar to the OGBG-PCBA dataset (Hu et al. 2020b), but instead of being limited to 128 assays and 437k molecules, it comprises 1,328 assays and 1.56M molecules. This dataset is very interesting for pre-training molecular models since it contains information about molecule's 's behavior in various settings relevant to biochemists, with evidence that it improves binding predictions (Laufkötter et al., 2019). Analogous to the gene expression, we obtain a bio- essay-expesssion of each molecule. To gather the data, we have looped over the PubChem index of bioassays (Kim et al., 2016) and collected every dataset such that it contains more than 6,000 molecules annotated with either "Active" or "Inactive" and least 10 of each. Then, we converted all the molecular IDs to canonical SMILES and used that to merge all of the bioassays into one dataset.
PCQM4M G25_N4 N4 This dataset comes from the same data source as the famously known for being part of the OGB large-scale challenge (Hu et al.) 202 and being one of the only graph datasets where pure Transformers have proven successful (Luo et al.) 2022: Ying
dataset,
et al. 2021). The data source is the PubChemQC project (Nakata & Shimazaki, (017)) that computed
DFT properties on the mergy-minimized conformation of 3.8M small molecules from PubChem. In contrast to the OGB challenge,
aim to provide enough data for pre- pre-traiming graph ML models, so we do not limit ourselves to the HOMO-LUMO gap prediction (Hu et al. 20216). Instead, we gather properties directly given by the DFT calculations (e.g., energies) and compute other 3ID descriptors from the conformation (e.g., inertia, the plane of best fit). We also gather node-level properties, in this case the Mulliken and Lowdin charges at each atom. Furthermore, about half of the molecules have time-dependent DFT calculations to help inform about the molecule's excited state. Looking forward, we plan on adding edge-level tasks enable the prediction of bond properties, such as their lengths and the gradient of the charges.
2.3 ULTRALARGE DATASET
PM6_83M This dataset is similar to PCQM4M and comes from the same PubChemQC project. However, it uses the PM6 semi-empirical computation of the quantum properties, which is orders of magnitude faster than DFT computation at the expense of lower accuracy (Nakata & Shimazaki
2017 Nakata et al. 2020).
This dataset covers 83M unique molecules, 62 graph-level tasks, and 7 node-level tasks. To our knowledge, this is the largest dataset available for training 2D-GNNs regarding the number of unique molecules. The various tasks come from four different molecular states, namely "SO" for the ground state, "TO" for the lowest energy triplet excited state, "cation" for the positively charged state, and "anion" for the negatively charged state. There are 221M PM6 computations (Nakata et al. 2020).
3 THE GRAPHIUM LIBRARY
Drug discovery is field with rich, multi-faceted datasets. Data is often inherently multi-tasked, containing label information at various levels, from nodes and edges to pairs of nodes and entire graphs. Hence, training models on such complex datasets requires strategy for jointly learning on multiple task levels, presenting a unique challenge, To leverage the rich information contained in these dataset mixes and to benefit from their scale, dedicated software is necessary to facilitate efficient training in multi-tas, task fashion. To this end, we introduce the Gr Graphium im. library, designed explicitly for lsage-scale multi-task task and multi-level ML models within the molecular domain.
Key features of the Graphi um library are summarised here, with further details to be found in Appendix [. with the relevant sections indicated in the text. The entire training procedure is defindd in modular configuration file, for example this allows users to swap out architecture, datasets or metrics for quick experimentation. The core novel feature of the Graphium um library is the multi-level multi-task learning that facilitates training model over multiple datasets with disparate and sparse labels through modular dataloading, task head, and loss functions. More details can be found in Appendix E. The flexible and modular modeling is detailed in Appendix E.22 Positional encoding methods, which provide node level information about location within the subgraph, are integral to many advanced molecular models. Key methods, including random walk and laplacian eigenvectors,

---

Published as a conference paper at ICLR 2024
provided in Graphium, this is detailed in Appendix E.3 Label normalisation (Appendix E.4). ranked classification loss (Appendix E5.)) and the handling of missing data (Appendix (.6) are important details of the library that facilitate the combination of multiple sources of data across range of tasks and the inevitable sparsity in the dataset. To make training large models resource efficient much attention is paid to the possibility of tuning model parameters on small models and applying these to larger models. One such approach called μP introduced in Yang et al. (2022) is included in Graphi um to reduce the cost of hyper-parameter tuning. More details are given in Appendix E.7 The library supports CPU, GPU and IPU [S. hardware to accelerate training. Further library optimisations are detailed in Appendix
EXPERIMENTS ON BASELINE MODELS
To demonstrate the capabilities of the Graphium um library in a multi-task setting with thousands of labels, a set of standard baselines were run with simple hyperparameter sweeps using 3 popular GNNs, namely GCN (Kipf & Welling 2017). GIN (Xu et 2019), and GINE (Hu et al. 2020a). A basic hyper-parameter sweep conducted for each model and multiple random seeds used for initialisation to provide a performance baseline for future experiments to be evaluated against. In Appendix H we report scaling law experiments on proposed datasets and observed significantly improved results as models include more parameters.
In addition to the node features (and edge features in the case of GINE) provided with the datasets, use (the first 8) eigenvalues and eigenvectors of the graph Laplacian as node-level PEs. Combining eigenvalues and eigenvectors and embedding them using an MLP offers PE informed by the graph's spectral properties (Kreuzer et al. 2021). We further use random walk return probabilities (i.e., the probability of a random walk starting a specific node and returning to the said node after k e N steps), again, in combination with an MLP encoder (Dwivedi et al. 2021). This can indicate the presence of cycles within molecules, which can be informative in the context of some molecular tasks, e.g., the regression of the solubility (logP) score on the ZINC dataset (Dwivedi et al. 2020). We finally compare the results from the multi-task model to a single-task model. Details of hardware and computational resources can be found in Appendix G
4.1 TOYMIX
Experimental setting. For all baseline models, the GNN module uses 4 GNN layers with a hidden size chosen to approximately balance the number of trainable parameters across models (reported in Tab. 2.. For TOYMIX, the models have approximately 150k parameters. We train a model on each dataset individually and the combined versions, with the results reported respectively in columns Single Dataset and Mulri Dataset. For these cases we fix both the learning rate and total number of epochs (300) used. Importantly this means that the total number of training steps will be larger for the combined dataset since there are more unique molecules.
Losses and metries computation. The multi-task loss is built by combining the loss function for each task. For the regression tasks, the models are trained with Mean Average Error (MAE) and evaluated using the Pearson correlation coefficient and the coefficient of determination R2, For the classification tasks, the loss is the binary cross entropy (BCE) and the metrics are area under the receiver operating characteristic curve (AUROC) and the average precision (AP) for the precision-recall curve. The MAE and BCE are computed on the flattened tensor, but the other metrics are computed per label and then averaged. All "NaNs" are filtered out from the computation. Also, we calculate the Mean Average Error (MAE) on regression tasks using the normalized outputs and labels so the results are not skewed by the magnitude of the labels themselves. Normalisation constants can be found with the dataset and future users should use these values to ensure results can be compared.
Comparability to results previous works. It should be noted that for the ToYMIX dataset, while the component datasets are well studied, the results present here are not directly comparable to prior work. Here QM9 has an extended list of 19 labels, that are normalized and aggregated, compared to 12 non non-normalized labels typically reported in prior work. ZINC 1 2K also uses 3 labels compared to the standard single label. Finally Tox21 uses a simple random split data split compared to the scaffold split used in prior literature.
7

---

Published as a conference paper at ICLR 2024
Table 2: Results for GNN baselines on the proposed TOYMIX dataset. We report performance metries on the test set (mean + std over 3 seeds) per dataset contained in Tox and for the dataset overall. The best score for each metric per dataset across all three models is in marked green, with the best result per model marked in orange.
Single Dataset Pearson f
Multi Dataset
Dataset
Model
MAE
R²
MAE J
Pearson 1
R²
GCN
.102 .003
.958 + 0007
.900
+ .002
.119 + 01
.955 # 001
.915 # .001
QM9
GIN
.0776 1 .006
.959 + 0002
.922
.0004
.117 +
.950 : 002
.908
= .003
GINE
.0959 +: .0022
.955
.918
+ .004
.102 +
.956 .009
.918
/002
GCN
.348 m
.941 + 002
.863 01
.226 #
.973 + .0055
.940
0013
ZINCI2K GINE
GIN
.303 + o00
.950 + 003
.889
+ 003
.189 .004
.978 + 2006
.953
.002
.266 + 0
.911 + 003
.915 + 01
.147 .009
.987 1 0011
.91
.003
BCE 4
AUROC
AP
BCE
AUROC 1
AP
GCN
.202 + .08
.773 + 006
.334 + 4.03
.176 + .001
.850 # .006
.446
.01
Tox21
GIN
.200 1 002
.789 + 009
.350 1.01
.176
.001
.841
1 .005
.454
.009
GINE
.201 + 007
.783 + 007
.345 + 02
.177 + .008
.836
+ .004
.455
.008
Results discussion. TOYMIX results can found Table We compare 3 different models and find that across all of the different settings, GINE outperforms GIN, which in turn outperforms GCN. To investigate the impact of using the union of multiple datasets we compare the performance for single dataset training with those multi dataset context. Interestingly we find that for the smaller ZINC12K and Tox21 datasets there is a clear advantage to also training them with the additional data. This is striking these tasks not highly related and the chemical spaces don't 't intersect much. This highlights that learning in unison with other molecular tasks alone may be beneficial, especially knowing that quantum data like QM9 and chemoinformatics data like ZINC12K cheaper to obtain than Tox21 I or other bioassay data. For the larger dataset QM9, which takes up nearly 87% of the molecules, however, we do see minor reduction in the performance in the multi dataset context, which suggests that there is some degree of compromise when using this setup.
much
4.2 LARGEMIX
Experimental setting. In case of LARGEMIX, we provide results for the same GNN baselines with 4 GNN layers and model sizes between 4M and 6M parameter. We present results of models trained for 200 epochs and average across 3 seeds. The remaining setup the same as for TOYMIx (Section 4. I reporting metrics on the Single Dataset and Multi Dataset are the same.
On the L100 datasets VCAP and MCF7, similar to TOYMIX, we observe clear advantage of our Multi Dataset approach compared to training on a single dataset. While not the case for the significantly larger PCQM4M and PCBA_ 1328 datasets. This again indicates the smaller datasets benefit more from the multi- dataset training approach. We believe that this is due to underfitting, as e.g., for PCBA_1325 a 10M parameter model is used train on than 3000 tasks jointly. Indeed, Table 7 observe that as we increase the model size, the results for PCQM4M_MM N4 and PCBA 1328 significantly improve.
Results discussion. LARGEMIX results can be found Table 3
4.3 ULTRALARGE
Experimental setting, For ULTRALARGE, we provide results for the same GNN baselines as for LARGEMIX. Fach model is trained for 50 epochs and results are averaged over 3 seeds. The remaining setup is the same as for ToYMIx (Section T.T), reporting metries on the Single Dataset and Multi Dataset using the same performance metrics. We further use the same models (in terms of size) as used for LARGEMIX. However due to resource constraints we have trained only on a 5% subset of the data.
Results discussion. ULTRALARGE results can be found in Table Interestingly, on both graph- and level tasks we observe that there is no advantage of multi-tasking in terms of performance. However, we expect the major benefits of this dataset to be found when used at scale and when either fine-tuning other quantum tasks or when used in conjunction with more diverse array of datasets.
8

---

Published as a conference paper at ICLR 2024
Table 3: Results for GNN baselines on the proposed LARGEMIX dataset. We report performance metrics on the test set (mean + std over 3 seeds) per dataset contained in LARGEMIX and for the dataset overall. The best score for each metric per dataset across all three models is in marked green, with the best result per model marked in orange.
Single Dataset
Multi Dataset
Dataset
Model
MAE +
Pcarson
R²
MAE
Pearson t
R²
GCN
.2362.:
.8781 .005
.7803= .0066
.2458
.8701 .002
.7678=
PCOM4M_G22 G25
GIN
.2270=
88854
.7912= .0066
.2352+
.8802+
2007
.7827=
GINE
.2223
.8874
.7949
.2315
.8823
.0022
.7864
GCN
.2080+ 0003
.5497-
.010
.2942
.0077
.2040
.0001
.4796+
.0066
.2185
.0002
PCQM4M_N4
GIN
.1912 .0027
.6138:
.088
.3688 .0166
.1966#
0003
.5198=
.0088
.2602
0012
GINE
.1910 .0001
.6127:1 0003
.3666
.0088
.19411 .003
.53031
.0033
.2701
.004
CE 4
AUROC1 1
AP
CE 4
AUROCT
AP
GCN
.0166
0000
.7960
1. .0000
.3368
.002
.0349
.002
.7661
.0011
.2527
.041
(322_1328
GIN
.0324
.7941
+ .3328
.0342+
.7747+
.2650=
GINE
.0320
.7944
.3337
.027
.0341
.7737=
.2611
GCN
.1900 = 0002
.5788
1 .004
.3708 1: 0007
.1872
.002
.6362.t
.0012
.4022.:
.0088
L1000_VCAM VCAP
GIN
.1909 = 0005
.5734
+ 0009
.3731
+ 0014
.1870
.001
.6351#
.0014
.4062
.0011
GINE
.1077
0006
.5708
1: .0099
.3705
0015
.1862.
0007
.6398
.0043
.4068
.0033
GCN
.1699 J003
.6123
.3866
.18634
.0011
.64014 + .0021
.4194
- .0044
L1000_MCF7
GIN
.1862 1 0003
.62022t
.091
.3876
0017
.1874
.013
.636771
=1 .0066
.4198 = .0336
GINE
.1856 1: .005
.6166
1 0017
.3892
0035
.1873a
.009
.6347
1: .048
.4177
0024
Table 4: Results for GNN baselines on 5% subset of the proposed ULTRALARGE dataset. We report performance metries the test set (mean + std over 3 seeds) per dataset contained in ULTRALARGE and for the dataset overall. The best score for each metric per dataset across all three models is in marked green, with the best result per model marked in orange.
Single Dataset
Multi Dataset
Datasct
Model
MAE 4
Pearson
R²
MAE
Pcarson
R²
GCN
.2606:: .0111
.9004:
.0044
.7997=
.009
.2625 0011
.8993 .005
.7977: 1: .008
PM6_83M 83M__662
GIN
.2541+ .017
.9053
.003
.8069_
.0028
.25784
2014
.8981
.007
.8031
0007
GINE
.2538+ .0006
.9057:
.088
.8078
.0013
.2575+
0009
.9043+
.0022
.8052+ 2003
GCN
.8033 .001
.3372:
.0033
.1188
.59711
.002
3164
.1019
PM6_ 23__83__N7
GIN
.5729 .0022
.3480
.003
.1271=
.0003
.5831+
.001
.3314=
2004
.1 .114 2008
GINE
.5721 1: .004
.3484::
.0066
.1262
.0006
.5839.
.004
.3294
.0033
.1108 2001
5 CONCLUSION
In this work, we curated a novel collection of molecular datasets with an unprecedented number of data points for supervised learning, thereby significantly enhancing the resources available for research in the field of drug discovery. These datasets are naturally multi-task and multi-level, posing unique challenges for machine learning models. To facilitate the study of these complex datasets, we presented the Graph i um library, framework designed to process and efficiently load large-scale molecular data and perform multi-task learning across various task levels. The library takes advantage of the unique features and properties presented by our datasets and build towards the training of large foundational model for molecular learning.
Furthermore, we present number of baseline results performance on small biological datasets benefits from being trained on large amounts of quantum data. We hypothesise that this improved generalization will help foundation models perform well when fine-tune tunet to the low-resource downstream tasks commonly encountered Drug Discovery.
proposed datasets and show that

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY
In line with ICLR's commitment to promoting reproducibility in research, we have taken diligent steps to ensure the reproducibility of our work presented in this paper. The Graphi library code can be accessed on Github https: httpps///gis //tp::////ttbtt com/ dat dam I 1 grgggrppiiig Sipppiiun um This code allows fellow researchers to reproduce and build upon our experimental results.
All datasets from TOYMIX, LARGEMIX and ultra-large can be found below. Datasets are provided either as .csv for smaller datasets or parquet for larger datasets. Splits are provided as .pt files, meant to be read via torch. load, and containing a dictionary with keys "train", "val", "test" and "test-seen", each containing list of indices corresponding the split.
Download links for the ToyMix. LargeMix and Ulraaarge datasets are found in Zenodo and Part 2/ttts::/zeeeee org
Part T/tt /t:httts:///eeooo fecords/8370547
ETHICS STATEMENT
We confirm that we adhere to the ICLR Code of Ethics as stated here This research primarily involves the analysis and prediction of molecular data using machine learning techniques. We have taken ethical considerations into account at various stages of our work. Our experiments do not involve human subjects, and we have not accessed any sensitive private information. The licenses for the datasets contributed in this work are discussed in Appendix B.2
ACKNOWLEDGMENTS
This work was conducted by Shenyang Huang, Gabriela Moisescu- Moisescu-Pareja, Oleksandr Dymov, and Frederik Wenkel during internship Valence Labs. Shenyang Huang is funded by NSERC PGS D Award and FRQNT Doctoral Award. Guillaume Rabusseau and Reihaneh Rabbany are supported by the CIFAR AI chair program. Frederik Wenkel is partially funded by the Fin-ML CREATE graduate studies scholarship for PhD, the J.A. DeSève scholarship for PhD and Guy Wolf's research funds. Guy Wolf is funded by IVADO (Institut de valorisation des données) grant FRQNT (Fonds de recherche du Québec Nature et technologies) grant 299376, NIH (National Institutes of Health) NIGMS grant RO1GM135929 and the Canada CIFAR AI Chair. Christopher Morris and Luis Müller are partially funded by a DFG Emmy Noether grant (468502433) and RWTH Junior Principal Investigator Fellowship under Germany's s Excellence Strategy. Michal Koziarski is funded by the Genetech-Mila partnership. Jian Tang is supported by Twitter, Intel, the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, Samsung Electronics Co., Ltd., Amazon Faculty Research Award, Tencent AI Lab Rhino-Bini Gift Fund, an NRC Collaborative R&D Project (AI4D-CORE-06 as well as the IVADO Fundamental Research Project grant PRF-2019-3588 139727. The content provided here is solely the responsibility of the authors and does not necessarily represent the official views the funding agencies.
10

---

Published as a conference paper at ICLR 2024
REFERENCES
Simon Axelrod and Rafael Gomez-Bombarelli, Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
Hyye Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35: 1423-11436, 2022.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro
Vinícius Flores Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Çaglar Gülçehre, H. Francis Song, Andrew J. Ballard, Justin Gilmer, George E. Dahl, Ashish Vaswani, Kelsey R. Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matthew M. Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. CoRR, (bbb/1806.0121,, 2018.
URL http:httt///ttiteee org/ abs 1806. 01261
Dominique Beaini, Saro Passaro, Vincent Létourneau, Will Hamilton, Gabriele Corso, and Pietro Liò. Directional graph networks. In International Conference on Machine Learning, pp. 748-758.
PMLR, 2021.
Camille Bilodeau, Wengong Jin, Tommi Jaakkola, Regina Barzilay, and Klavs Jensen. Generative models for molecular discovery: Recent advances and challenges. Wiley Interdisciplinary Reviews: Computational Molecular Science, 12, 03 2022.
Cristian Bodnar, Fabrizio Frasca, Nina Otter, Yuguang Wang, Pietro Lio, Guido F Montufar, and Michael Bronstein. Weisfeiler and lehman go cellular: Cw networks. Advances in Neural Information Processing Systems, 34:2625-2644 2021.
Nathan Brown, Marco Fiscato, Marwin S. Segler, and Alain C. Vaucher. Guacamol: Benchmarking models for de novo molecular design. Journal of Chemical Information and Modeling, 59 1011099-1108, March 2019. ISSN 1549-9596, 1549-960X. doi: 1a.022//as..cim.800000
arXiv:1811.0962 Iphysics, q-bio].
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020a.
Tom B. Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss., Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020b.
Yukuo Cen, Zhenyu Hou, Yan Wang, Qibin Chen, Yizhen Luo, Zhongming Yu, Hengrui Zhang, Xingcheng Yao, Aohan Zeng. Shiguang Guo, Yuxiao Dong, Yang Yang, Peng Zhang, Guohao Dai, Wang, Chang Zhou, Hongxia Yang, and Tang. Cogdl: comprehensive library for graph deep learning. In Proceedings of the ACM Web Conference 2023 (WWW'23), 2023.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Doming Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 11((0):6059-00 007000996777...000 2021.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking, arXiv preprint
2022. https://github.
CSIRO's Data61.
Stellargraph machine learning library.
stell la wlwaaaggghsssscc 1argraph 2018.
11

---

Published as a conference paper at ICLR 2024
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations.
Vijay Prakash Dwivedi, Chaitanya K. Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. CoRR, abs/2003.00982, 2020. BtLhttps: pentts://arxxe //axivv org/ httg/bbb/20033. 00982
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. arXiv preprint arXiv2211000787 2021.
Peter Eastman, Pavan Kumar Behara, David L. Dotson, Raimondas Galvelis, John E Herr, Josh T Horton, Yuezhi Mao, John D Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. Spice, a dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop Representation Learning Graphs and Manifolds, 2019.
Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. International Conference on Learning Representations.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34: 6790-6802, 2021.
Jonathan Godwin Thomas Keck*, Peter Battaglia, Victor Bapst, Thomas Kipf, Yujia Li, Kimberly Stachenfeld, Petar Veličković, and Alvaro Ganzhezz Jraph: A library for graph neural networks in jax., 2020. URL http: htt///gtttubbbt com dewpmind//d/eee//ee jraph/
Hatem Helal, Jesun Firoz, Jenna Bilbrey, Mario Michael Krell, Tom Murray, Ang Li, Sotiris Xantheas, and Sutanay Choudhury, Extreme acceleration of graph neural network-based prediction models for quantum chemistry, 2022.
Shion Honda, Shoi Shi, and Hiroki R. Ueda. SMILES transformer: Pre-trained molecular fingerprint for low data drug discovery. CoRR, abs/1911.044338, 2019. URL [Rtp:ttt://arxiv.oo// /bs/ 1911 04738
Jun Hu, Shengsheng Qian, Quan Fang, Youze Wang, Quan Zhao, Huaiwen Zhang, and Changsheng Xu. Efficient graph deep learning in tensorflow with tf geometric. In Heng Tao Shen, Yueting Zhuang, John R. Smith, Yang Yang. Pablo Cesar, Florian Metze, and Balakrishnan Prabhakaran (eds.), MM '21: ACM Multimedia Conference, Virtual Event, China, October 20 24, 2021, pp. 3775-3778. ACM, 2021a. doi: 10.1145/37400553477222 URL /tttps://doi.oo otg.///0000000/
1145/3474083 3478322
W Hu, B Liu, J Gomes, M Zimik, P Liang, V Pande, and J Leskovee. Strategies for pre-training graph neural networks. In International Conference on Learning Representations (ICLR), 2020a.
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 Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, viriual, 2020b.
nwurips.....em httpipp..ccppppee httpp//pppp///2///aaaaaa. 1160d411a5c550 1tl.aaaccc722222335222aaaaa F550001112c5522222333222errrr0-0 Abst .tt...tttt..................m
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec, OGB-LSC: A large-scale challenge for machine learning on graphs. In Joaquin Vanschoren and Sai-Kit
Yeung (eds.), ). Proceedings of the Neural Information Processing Systems Track on Datasets and
Benchmarks 1, NewrIPS Datasets and Benchmarks 2021, December 2021, virtual, 2021b. URL
cc http:a/iip...cppppe///222.. http:///aaaaooooaaaaaaaaaaaaa4a. htterraccctoooo2h.tmmm..
12

---

Published as a conference paper at ICLR 2024
Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. arXiv preprint wxXi:2212.00954. 2021.
John J. Irwin, Khanh G. Tang, Jennifer Young, Chinzorig Dandarchuluun, Benjamin R. Wong, Munkhzul Khurelbaatar, Yurii S. Moroz, John Mayfield, and Roger A. Sayle. Zinc20-a free ultralarge-scale chemical database for ligand discovery. Journal of Chemical Information and Modeling, 60(12):6065-6077
December 2020. ISSN 1549-9596. doi: 1.102//acc.jeim.0000677.
Zhe Jia, Blake Tillman, Marco Maggioni, and Daniele Paolo Scarpazza. Dissecting the grapheore ipu architecture via microbenchmarking, 2019.
Kuzma Khrabrov, Ilya Shenbin, Alexander Ryabov, Artem Tsypin, Alexander Telepov, Anton Alekseev, Alexander Grishin, Pavel Strashnov, Petr Zhilyaev, Sergey Nikolenko, et al. nabladft: Large-scale conformational energy and hamiltonian prediction benchmark and dataset. Physical Chemistry Chemical Physics, 24(42):25533 2022.
Sunghwan Kim, Paul A Thiessen, Evan E Bolton, Jie Chen, Gang Fu, Asta Gindulyte, Lianyi Han, Jane He, Siqian He, Benjamin A Shoemaker, et al. Pubchem substance and compound databases. Nucleic acids research, 44(DI):D1022-D121 2016.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24- 26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL net /
Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34:21618-2166 2021.
Oliver Laufkötter, Noé Sturm, Jürgen Bajorath, Hongming Chen, and Ola Engkvist. Combining structural and bioactivity- fingerprints improves prediction performance and scaffold hopping capability. Journal of cheminformatics, 11:1-14, 2019.
Xinyuan Lin, Chi Xu, Zhaoping Xiong, Xinfeng Zhang, Ningxi Ni, Bolin Ni, Jianlong Chang, Ruiqing Pan, Zidong Wang. Fan Yu, et al. Pangu drug model: learn a molecule like human. bioRxiv, pp. 2022-03 03, 2022.
Meng Liu, Youzhi Luo, Limei Wang. Yaochen Xie, Hao Yuan, Shurui Gui, Haiyang Yu, Zhao Xu, Jingtun Zhang, Yi Liu, Keqiang Yan, Haoran Liu, Cong Fu, Bora M Oztekin, Xuan Zhang, and Shuiwang Ji. DIG: A turnkey library for diving into graph deep learning research. Journal of Machine Learning Research, 22(240):1-9, 2021. URL htth:t/::/ //rmmmm..... www.:.orgpppppppee pttp://ppprrs//221
21 21-0343. html
Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with aev3iainvariant denoising distance matching, 2023a.
Yiheng Liu, Tianle Han, Siyuan Ma, Jiayue Zhang, Yuanyuan Yang, Jiaming Tian, Hao He, Antong Li, Mengshen He, Zhengliang Liu, al. Summary of chatgpt/gpt-4 research and perspective towards the future of large language models. arXiv preprint arXi::2304.01852 2023b.
Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2D & 3D molecular data, arXiv:22 2022.
Dominic Masters, Josef Dean, Kerstin Klaser, Zhiyi Li, Sam Maddrell- Adam Sanders, Hatem Helal, Deniz Beker, Andrew Fitzgibbon, Shenyang Huang, et al. Gps++: Reviving the art of message passing for molecular property prediction. arXiv preprint arXiv:2302.02947. 2023.
Alexander Mathiasen, Hatem Helal, Kerstin Klaeser, Paul Balanca, Josef Dean, Carlo Luschi, Dominique Beaini, Andrew William Fitzgibbon, and Dominic Masters. Generating QM1b with PySCF ipu In Thirty-seventee Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023.
http://pppeeeeeeieeeeeeeee net. m d c cm0757o
13

---

Published as a conference paper at ICLR 2024
Andreas Mayr, Günter Klambauer, Thomas Unterthiner, Marvin Steijaert, Jörg K. Wegner, Hugo
Ceulemans, Djork-Amé Clevert, and Sepp Hochreiter. Large-scale comparison of machine learning methods for drug target prediction chembl. Chemical Science, 9044)5441-5451, June 2018.
ISSN 2041 2041-6539. doi: 10. 10.1039//880)))44..
Oscar Méndez-Lucio, Christos Nicolaou, and Berton Earnshaw. Mole: molecular foundation model for drug discovery. arXiv preprint wrXiv:221..02635 2022.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project:
large-scale first- prsst-priccipls electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 57(6): 1300-1308, 2017.
Maho Nakata, Tomomi Shimazaki, Masatomo Hashimoto, and Toshiyuki Maeda. Pubchemqe pm6: Data sets of 221 million molecules with optimized molecular geometries and electronic properties. Journal of Chemical Information and Modeling, 60 60(12):5891-889) 2020.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
R. Ramakrishnan, O. Dral, P., M. Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014. Nature.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35: 15:14501-145155 2022.
Bharath Ramsundar, Peter Eastman, Patrick Walters, Vijay Pande, Karl Leswing, and Zhenqin Wu. Deep Learning for the Life Sciences. O' O'Reilly Media. 2019, https: www. amazon. com/ Deep- Sear-iLeaning--- Life Sciences
Lars Ruddigkeit, Ruud van Deursen, Lorenz C. Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb- Journal of Chemical Information and Modeling, 52(11 1::264 2875, 2012. doi: 10.102//ci30441d.
Johannes Schimunek, Philipp Seidl, Lukas Friedrich, Daniel Kuhn, Friedrich Rippmann, Sepp Hochreiter, and Günter Klambauer. Context-enriched molecule representations improve few-shot drug discovery. arXiv preprint arXiv:2305.0948/, 2023.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, ACL 2016, August 7-12, 2016, Berlin, Germany, Volume 1: Long Papers, 2016.
Till Siebenmorgen, Filipe Menezes, Sabrina Benassou, Erinc Merdivan, Stefan Kesselheim, Marie Piraud, Fabian Theis, Michael Sattler, and Grzegorz M Popowicz. Misato-machine learning dataset of protein-ligand complexes for structure-bart drug discovery. bioRxiv, pp. 2023-05. 2023.
Megan Stanley, John F Bronskill, Krzysztof Maziarz, Hubert Misztela, Jessica Lanini, Marwin Segler, Nadine Schneider, and Mare Brockschmidt. Fs-mol: A few-shot learning dataset of molecules. In Thirty-fitth Conference on Neural Information Processing Systems Datasets and Benchmarks Truck (Round 2), 2021.
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, pp. 20479-20502. PMLR, 2022.
Aravind Subramanian, Rajiv Narayan, Steven M Corsello, David D Peck, Ted E Natoli, Xiaodong Lu, Joshua Gould, John F Davis, Andrew A Tubelli, Jacob Asiedu, et al. A next generation connectivity map: L.1000 platform and the first 1,000,000 profiles. Cell, (71
2017.
Ruoxi Sun, Hanjun Dai, and Adams Wei Yu. Does gnn pretraining help molecular representation? Advances in Neural Information Processing Systems, 35:12096-1229 11109, 2022.
14

---

Published as a conference paper at ICLR 2024
Susheel Suresh, Pan Li, Cong Hao, and Jennifer Neville. Adversarial graph augmentation to improve graph contrastive learning. In Mare' Aurelio Ranzato, Alima Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, I, NeurlPS 2021, December 6-14, 2021, virtual, pp. 15920-159333 2021. URL htt.Ltttttt:://rrceeeeeeet
neur ps CC/ Abstract -tmm
Raymond R Tice, Christopher P Austin, Robert Kavlock, and John R Bucher. Improving the human hazard characterization of chemicals: tox21 update. Environmental health perspectives, 121(7): 756-765, 2013.
R. Todeschini and P. Gramatica. The whim theory: New 3d molecular descriptors for qsar in environmental modelling. SAR and QSAR in Environmental Research, 7(1-4):89-113 1997 doi: 10. 1080/1062933970003391 URL https::/dbi..org/// htt..//../........00000002333330
Researcc
Derek van Tilborg, Alisa Alenicheva, and Francesca Grisoni. Exposing the limitations of molecular machine learning with activity cliffs. Journal of Chemical Information and Modeling, 62:5938-
5951, 2022. doi: 1a022//as..eii2220007...
Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis, Jinyang Li, and Zheng Zhang. Deep graph library: A graph-centric, highly-pene highly-performant package for graph neural networks. arXiv preprint ar arXhv:190.001315. 2019.
Zichen Wang, Neil R Clark, and Avi Ma' My'ayan. Drug-induced adverse events prediction with the lines 11000 data. Bioinformatics, 32(15):2338 2345, 2016.
Robin Winter, Floriane Montanari, Frank Noé, and Djork-Amé Clevert. Learning continuous and data-driven molecular descriptors by translating equivalent chemical representations, Chemical Science, 10(6):1692-17000 February 2019. ISSN 2041-6539, doi: 10. 10.103//C8C041155.
Fang Wu, Shuting Jin, Yinghui Jiang, Xurui Jin, Bowen Tang. Zhangming Niu, Xiangrong Liu, Qiang Zhang. Xiangxiang Zeng. and Stan Li. Pre-training of equivariant graph matching networks with conformation flexibility for drug binding, Advanced Science, 9(33):2003796. 2022.
Z. Wu, B. Ramsundar, N. Feinberg, J. Gomes, C. Geniesse, S. Pappu. K. Leswing, and V. Pande. MoleculeNet: A benchmark for molecular machine learning. Chemical Science, 9:513-530, 2018.
Jun Xia, Yanqiao Zhu, Yuanqi Du, Yue Liu, and Stan 7. Li. A systematic survey of chemical pre-trained models. IJCAI, 2023.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URI htt.:httss///cm ftttt::://openrevviee ew. net./ forum?
id (d-ryGa6 Km
Ge Yang. Edward Hu. Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 17:17774-77097, 2021.
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: tuning large neural networks via zero-shot hyperparameter transfer. CORR, abs/2203.03446 2022. doi: 10.48550/arXiv.2200..
03466. URL o.t.o///...taaaooooo.oooo./.ooa.
Xingyi Yang, Xuehai He, Yuxiao Liang, Yue Yang, Shanghang Zhang, and Pengtao Xie. Trans- fer learning or self-supervised learning? a tale of two pretraining paradigms. arXiv preprint arXiv:2007.04234,, 2020.
Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances Neural Information Processing Systems, 34:28877-28888 2021.
15

---

Published as a conference paper at ICLR 2024
Jiaxuan You, Rex Ying, and Jure Leskovec. Design space for graph neural networks. In NeurIPS, 2020.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter W. Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. CoRR, abs/2206.00133. 2022. doi: 10.48550/arXiv 2206.00133. URL htttss//dddd owg. //c 48w.0//rxxxii..... 2206. 00133
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. 2023,
Zhaocheng Zhu, Chence Shi, Zuobai Zhang, Shengehao Liu, Minghao Xu, Xinyu Yuan, Yangtian Zhang. Junkun Chen, Huiyu Cai, Jiarui Lu, Chang Ma, Runcheng Liu, Louis-Pascal A. C. Xhon- neux, Meng Qu and Jian Tang. Torchdrug: A powerful and flexible machine learning platform for drug discovery. CoRR, abs/2202.08320,, 2022. URL StLtttps: : /ttr//aavvvvv// argg/bs//200 08320
16

---

Published as a conference paper at ICLR 2024
A LIMITATIONS, FUTURE WORK AND SOCIETAL IMPACTS
Limitations and Future Work Our proposed datasets currently do not contain edge- or nodepair- level tasks, which we plan include in the future. Also, benchmark results limited few baseline models, but we will continue to run experiments and ablations and eventually include more positional encodings and graph Transformers. Further, large-scale multi-task learning on millions of molecular data may pose a limitation in situations where extensive computing resources are unavailable. We will continue work on optimization, including chip-specific accelerations, to help reduce the time and cost of training. On that same note, the ultra-large PM6 83M dataset is so large that using the LargeMix might be challenging since in the current Graphi um framework, 95% of the molecules in any batch will come from that specific dataset, and under-sampling might be required. Finally, since the capability of a model trained on multiple tasks to adapt and perform well on new tasks a crucial indicator of its robustness, we aim investigate our multi-task framework's ability to transfer to new tasks.
Societal Impact Releasing our datasets and the Graphi um library may not have immediate direct societal impacts. However, it is crucial to acknowledge the potential implications that arise when providing access to a foundation model for molecular graphs. One concerning possibility is the misuse of this technology for the development of chemical weapons, toxins, or unregulated drugs. To address these potential risks, are committed implementing robust mitigation strategies. Central to our approach is the active promotion of beneficial applications, particularly in the fields of material and drug discovery. By highlighting the positive utilization of this technology, we aim channel its potential toward scientific advancements that contribute to societal well-being.
B
DATASET DOCUMENTATION AND INTENDED USES
B. INTENDED USE
The datasets are intended to be used in academic setting for training molecular GNNs with orders of magnitude more parameters than current large models. Further, the TOYMIX and LARGEMIX datasets are intended to be used in multi-takk setting, meaning that single model should be trained to predict them simultancously. We invite users to use the Graphi um library and contribute to it to simplify the development state-of-the-ar GNNs on multi-task and multi-label data.
B.2 DATASET LICENSES
The licenses of the existing datasets used in this work is follows:
QM9: Creative Commons Attribution Non-Commercial Share 4.0
Tox21: Attribution 4.0 International (CC BY 4.0)
Zinc12k: MIT license
L1000: free to use by all
PCQM4M and PM66_3MM Creative Commons Attribution 4.0 International license.
PCBA 1328: Open access
All datasets provided in this work are licensed under the Attribution Non-Commercial ShareAlike 4.0 International (CC BY-NC-SA 4.0) license We chose this license because some of the original datasets have this license and we provide our datasets with the same level of access.
B.3 MOLECULAR PROPERTIES
For the PM6 83M dataset, we added many chemical properties using rdkit, which predicted together with the quantum properties. This enables the trained model to learn a shared embedding between chemical and quantum properties.
The computed chemical properties are:
MW: molecular weight.
17

---

Published as a conference paper at ICLR 2024
fsp3: fraction of sp3 carbon atoms
TPSA: topological polar surface area
QED: quantitative estimation of drug-likeness
clogp: partition coefficient for n-octanol/water (measure of lipophilicity)
SAS: synthetic accessibility score
n p'lipinskk_hae hba: number of hydrogen bond acceptors
n pillpisskk__pee hbd: number of hydrogen bond donors
n_rings: number of ring structures
n hetero_atoms: number heteroatoms (non hydrogen or carbon atoms)
n heavy_atoms: number of non-hydrogen atoms
rotatable_ bonds: number of rotatable bonds
radical_electrons: number of radical electrons
Seiphaticccarbocyeheee number of aliphatic carbocycles
allphaticcheterooyry heterocyles: number of aliphatic heterocycles
aliphatic_rings: number of aliphatic rings
aromatic_carbocyylee number of aromatic carbocycles
number of aromatic heterocycles
n coomatic_ringss number of aromatic rings
n_ ceturated_carbocyyee number of saturated carbocycles
n saturated heterrccleeeee number of saturated heterocycles
n saturated rings: number of saturated rings
C VISUALIZING THE DATASETS
C.1 A FEW MOLECULES
We randomly sample 5 molecules from each dataset and visualize them in Figure 2 The PM6 83M dataset omitted since it includes the PCQM4M_G25 N4 and the PCBA_1328.
C.2 MIX VISUALISATION
The chemical spaces of the dataset mixes are visualised in Figure showing the overlap between dataset components. In the TOYMIX, observe that ZINC12K and Tox21 share part of the chemical space, although Tox21 is more diverse, but QM9 is very different from both of them. In the LARGEMIX. we observe that there is strong intersection between the datasets the center of the plot, but that the MCF7 and VCAP, despite being smaller, cover larger regions of the chemical space, justifying their usefulness for pre-training foundation models of chemistry.
18

---

Published as a conference paper at ICLR 2024
QM9
goot
Tox21
dios tholi giu ooog
ZINC12K
youd &r
L1000_VCAR '0
100
L1000_MCF7
PCBA_1328
#CQM4M_G25 G25_N4
Figure Visualizing random samples from the datasets.
DVS
MCF
PORA
2IN
POOMTA
VCCP
UMMP mwstss www.heerrss0.. min. d66-0.86
(a) MoxMix
(b) LargeMix
Figure 3: UMAP visualizations for the dataset mixes, computed with the ECFP molecular fingerprints, showing the chemical spaces of the datasets.
19

---

Published as a conference paper at ICLR 2024
Mix
Graph labels
Node labels
Total
TOYMIX
2.6M
2.6M
LARGEMIX
344M
197.7M
541.7M
ULTRALARGE
4B
8.5B
12.5B
Total
4.3B
8.7B
13.04B
Table 6: Overview over the number of graph and node labels mix.
C.3 DATASET STATISTICS
Table 5: General statistics about the datasets, where G. / N. respectively indicate graph-level node- level tasks, and where the sparsity indicates the % of missing data. atoms and bonds denote respectively the heavy atoms, and the bonds between 2 heavy atoms. For data type, Q. indicates it is Quantum, C. means it is Computed, B. means it Bio-assays and means it is Transcriptomics. For task, C., R., and RC. denote classification, regression and ranked classification respectively.
Mix
TIYMIX ZINCI2K
LARGEMIX
ULTRALARGE PM6_83M
Dataset
QM9
Tox21
PCQMAM_G25_NG
PCBA_
L1000_vCAR VCAP L1000_MCH7
Type
C.
Task
R.
R.
RC
RC
R.
mols
133.8k
12k
7.8k
3.8M
1.5M
15.2k
11.6k
83.7M
labels
1,328
G. data
2,5M
36k
74.9k
93.4M
224, 3M
14.8k
11.3k
4,00B
% G. sparsity
20.2
1.9
89.2
0
22.9
labels
N. data
197.7M
8.51B
% N. sparsity
7.9
40.8
atoms
278k
h. bonds 1.26M
299k
151k
55.2M
44.8M
550k
395k
2.20B
We analyze the different datasets by using histogram over relevant chemical properties in Figure 4 There, we observe that the quantum datasets (QM9 and PCQM4M) contain much smaller molecules than the other datasets. This could hinder their ability to produce an embedding that generalizes well across the chemical space. However, the semi-empirical dataset PM6 83M, has very similar distribution to biologically relevant datasets, which makes it ideal for pre-training foundation model.
D DEEPER DIVE INTO THE LARGEMIX AND ULTRALARGE DATASETS
In this section, we into further details to explain what labels are present in the datasets, and how they compare to the current literature.
D.1 PCQM4M_ 10AAM_GG55MM0_ G25_N4
As mentioned in section 2.2 this dataset comes from the same data source as the OGBG-PCQM4M, namely the PubChemQC project (Hu al. 202Tb), meaning that it comprises of the same molecules.
However, this dataset differs from DGBG-PCQM4M since additional properties are included, and the 3D descriptors are provided as labels. This implies that we can't 't use 3D structure as input, as done by OGB, since it is contained in the labels and would lead to data leaks. Further, the OGB challenge only provides labels for HOMO-LUMO gap predictions alongside the 3D structures for the training set. The reasoning behind OGB's choice of including the 3D conformation only in the training set comes from their desire that model learns to implicitly embed when predicting its associated HOMO-LUMO gap. However, since work in multi-task setting, it is harder for us to use the same technique. For example, suppose that we are also predicting the binding affinity as another task, then the 3D conformation responsible for this affinity not necessarily the same the one provided in the PCQM4M dataset. Therefore, our choice including the 3D conformation the labels will allow the model to better associate different conformation to different tasks.
20

---

Published as a conference paper at ICLR 2024
TPSA
GE3
38000
190
109
1800
200
106
1m5
2000
10000
140000
19000
2000
8000
10000
TPSAA
GEO
300
TPSA
CC0
000
1500
0.4
08
1.
TPSA
GF0
10:00
000
1.00
100
200
1500
Figure 4: Histograms of the datasets for various properties relevant to chemists. These include molecular weights (MW), topological polar surface area (TPSA), quantitative estimation of drug- likeness (QED), and synthetic accessibility score (SAS). MW is the sum total of the atomic weights for all atoms in a molecule. TPSA is the surface sum over all polar atoms (like oxygen and nitrogen) and estimates the surface area of a molecule that is available for interactions with polar groups in biological systems. QED measures how drug-like a given compound is and is based on various molecular properties like MW, lipophilicity, hydrogen bond donors/acceptors TPSA and the presence of certain functional groups. SAS a metric to assess the feasability of synthesizing a given compound and is calculated based several descriptors, such size, complexity and stereochemistry. For the PM6 dataset, we calculated and plotted the descriptors for 100 molecules.
21

---

Published as a conference paper at ICLR 2024
The DFT properties that predict at the graph level are the following:
Alpha HOMO
Alpha HOMO-LUMO gap
Beta HOMO
Beta HOMO
Total energy
should be noted here that only the Alpha HOMO-LUMO gap is provided in the original OGBG- PCQM4M dataset. For most molecules, the number of labels are duplicated for time-dependent DFT (TDDFT), which allows to also model the excitation states of the molecule.
Using the 3D conformation optimized by the DFT, we further computed the following 3D descriptors as labels:
Spherocity
Plane of best fit
Principal length
Mass inertia
Valence inertia
Mulliken charges inertia *
Lowdin charges inertia
Values denoted with "3D" means that the descriptors are computed alongside the main axes of the molecule. The descriptors are inspired by existing 3D descriptors such as WHIM (Todeschini & Gramatica 1997].
Finally, the node property predictions are the atomic charges for both DFT and TDDFT when available.
Mulliken charges
Lowdin charges
D.2 L1000 VCAP AND MCF7
As discussed in section 2.2 the LINCS L,0000 a database of high-throughput transcriptomics that screened more than 30,000 perturbations set of 978 landmark genes (Subramanian al.) 2017) from multiple cell lines. VCAP and MCF7 are, respectively, prostate cancer and human breast cancer cell lines, with more details given below.
MCF7 is a breast cancer cell line possessing drug-discovery relevant characteristics of differentiated mammary epithelium, such as estrogen receptors and dome formation3. Isolated in 1970 from woman with metastatic adenocarcinoma of the breast, it became widely used as model for breast cancer research and drug testing.
VCAP is prostate cancer cell line exhibiting drug-discovery relevant features of advanced prostate cancer, such as androgen receptor signaling and bone metastasis. Isolated 1997 from man with hormone refractory prostate cancer, it became widely used as a model for prostate cancer research and drug testing.
Many perturbagens are used, but in our case, we focus only on small molecular perturbagens. Due to the high noise the assays, we convert the task multi-class classification, with one class having perturbation z-score < -2 and another > and the neutral class with perturbation in the range [_
There is lot of noise in the L1000 regression dataset which can make it difficult to give accurate predictions, therefore, the floating maiting-pointt dataset converted into a binary classification task. By applying a threshold, make the signal easier to distinguish from noise. While this can occasionally lead to mislabeled elements, our experience demonstrated that a classification task performed more reliably then regression task.
22

---

Published as a conference paper at ICLR 2024
D.3 PCBA 1328
This dataset is gathered in similar way to the OGB-PCBA (Hu et al. 2020b), by scraping some of the largest assays from Pubchem, but much larger scale.
Due to the unknown noise present in each assay, the variability between repeated experiments, and the difficulty associated with properly curating the datasets, we followed the same process as the OGB-PCBA data. Namely, we considered only the assays with author's 's labeling that mention "Active" for the positive class and "Inactive" for the negative class. All other labels were considered as missing values, since they are often about experimental errors.
As described in Appendix D.I, this dataset has several key differences to the original PBB8PMMMMM dataset. First, it covers 1328 assays, 10-fold increase, and 1.5M unique molecules, a 3-fold increase. This is due to our less stringent process of selecting the assays. Indeed, every assay prior to 2022 with more than 6000 datapoints and more than 20 points in the positive class is selected. This implies that other popular datasets such as Tox21 (Huang et al. (02I)) and OGB-HIV (Hu et al.) 2020b are also included in the dataset. Further, all 561 protein assays from the tigh-throughput fingerprints (Laufkötter et al. 2019) are present, with a prior demonstration the meaningfulness of their molecular representation.
We cannot name all the assays present since the process was automated. We further omitted assays from 2022 and 2023 to allow future work them for benchmarking purposes.
D.4 RM6_83M
As discussed in section 2.3 this dataset contains PM6 semi-empirical computation of quantum properties, a faster but less accurate computation than DFT's (Nakata & Shimazaki] 2017; Nakata et al. 2020). Covering 83M unique molecules, 62 graph-level tasks, and 7 node-level tasks, this is the largest dataset available for training 2D-GNNs regarding the number of unique molccules. In total, there are 221M PM6 computations Nakatta al. 2020).
Just like the PCQM4M G25 _N4 dataset contains an excited state, this one contains four different molecular states, namely "SO" for the ground state, "TO" for the lowest energy triplet excited state, "cation" for the positively charged state, and "anion" for the negatively charged state.
G25_N4
The quantum properties that predict the graph level are the following. For many molecules, the number of labels are multiplied by to account for the SO, TO, cation, and anion states of the molecule.
Alpha HOMO
Alpha HOMO-LUMO gap
Beta HOMO
Total energy
It should be noted here that only the Alpha HOMO-LUMO gap is provided in the original OGBG- PCQM4M dataset.
Using the 3D conformation optimized by PM6, we further computed the following 3D descriptors as labels. Values denoted with "3D" means that the descriptors are computed alongside the 3 main axes of the molecule. We compute less descriptors than the PCQM4M to avoid having too many properties, especially considering that there are now 4 excited states and a lot more molecules.
Spherocity
Plane best fit
Principal length *
Then, the node property predictions are the atomic charges and spins for all SO, TO, anion and cation states.
Mulliken charges
23

---

Published as a conference paper at ICLR 2024
Electronic spin
Finally, all 22 properties from B.3 are also added labels.
D.5 OTHER RELEVANT DATASETS IN THE LITERATURE
There are many other large datasets in the literature of molecular discovery. In this section, we give a brief description of a few of these datasets and how they are different or not suited to our current problem.
First, there are few large molecular datasets aimed towards material discovery. For these, both the chemical space and the set of tasks / molecular properties don't 't intersect. Such examples NablaDFT (Khrabrov et al. 2022 and OpenCatalyst (Chanussot et al. 2021), with the former being focused on trajectories.
Then, there are other datasets more related to drug discovery, but whose goals are to understand molecular dynamics within pockets MISATO (Siehenmorgen al., (023) to learn more accurate force-fields Spice (Eastman et al.) 2023). Both of these tasks are not suited for our current desire of training a large 2D-GNNs that work on the graph alone, although they could be suited for 3D-GNNs.
Finally, datasets such as GEOM (Axelrod & Gomez-Bombarehe 2022) were shown to be very useful for pre-training GNNs, however they require dedicated architectures and their performance remain limited (Stärk et al.
2022 Liu et al. 2023a).
D.6 RELATION BETWEEN THE DATASETS
There are very little direct relations between the datasets as they cover a wide variety of tasks. However, our empirical results show that running in a multi-task setting improves the results pared to single-task experiments. This is because there are some relationships between them, and understanding one can lead better understanding of the other.
First, the quantum simulation datasets, namely PCQM4M B00044__255_D0 G25_N4 and PM6 _8MM,,,,,222 contain information about the inner workings of molecules, how changing atoms affect the charges, the polarity, the 3D structure, the energy and excitation energies, the electronic densities, etc. All these are relevant when considering binding affinity protein, thus they are expected to help improve generalization on the binding assays that make up the majority of PCBA 1328.
We further argue that many ADME (absorption, distribution, metabolism, and excretion) properties depend on the physies and shape of the molecules, and they are very relevant to both binding and cell assays. Hence, learning on the quantum datasets likely to benefit the other ones.
Second, learning the binding affinity alongside other assays in PCBA_ 1328 can lead to an im- provement in transcriptomies predictions. Indeed, molecules mostly affect the cells via inhibition or activation of proteins, and this effect will often be visible via the changes in gene expression. Hence, we expect an improvement of the L 1000 predictions when combining to PCBA 1328, and by induction, to the other quantum datasets.
Third, learning on 1000 transcriptomic expected to lead to an improvement in the prediction of biological properties of molecules since the studied cell-lines are simpler models of the human biology and diseases. Hence, despite not providing such dataset in the current work, we expect that fine-tuning pre-trained models on such tasks will be beneficial.
E FURTHER DETAILS ON THE GRAPHION LIBRARY
The following section provides further details about the Graphium library.
E. MULTI-LEVEL, MULTI-TASK AND MULTI-LABEL. LEARNING IN GRAPHIUM
Graphium provides efficient mechanism for loading and processing multi-task datasets and handling tasks in parallel to enhance computational efficiency significantly. It supports multiple tasks, including node, edge, nodepair, and graph-level prediction tasks. Additionally, the library
24

---

Published as a conference paper at ICLR 2024
incorporates multi-level positional encodings; see Appendix E.3 capacity to capture and learn from the inherent hierarchical relationships within the molecular data (Rampäšek et al. 2022), providing unified platform for handling diverse data requirements. In Graphiuu, we propose a multi-task, multi-level model framework as illustrated in Figure []] After pooling, the architecture utilizes GNN layers to extract node, edge, nodepair, and graph-level Teatures. The resulting features are subsequently fed into task-level specific MLPs, which, in turn, are fed into task-specific heads, generating predictions at node, edge, nodepair, and graph-level tasks. This multi-level prediction strategy accounts for complex hierarchical relationships within molecular data, allowing the model to perform simultaneous learning tasks at different scales.
Such encodings enhance the model's
Multi-level Learning: In Graphium The term multi-level refers to model's capability to learn and output representations at various levels of the graph. Graphi um supports four levels: node, edge, graph, and node-pair. To simplify this, let's use molecular networks as an example.
On the node level, the task is to predit property of the individual nodes within a molecular graph. Each node often represents atom; hence, node-level task might involve learning representations of an atom's 's properties or its role within a particular molecular graph. The edge level focuses the relationship between connected nodes in graph. In molecular networks, the edge level aims to learn the embeddings of the type of bond (e.g. single, double or aromatic) or the relationship between atoms proteins.
On the graph level, a property of the entire graph is considered. molecular networks, this involves learning the representations of molecular properties, like its solubility stability.
Lastly, the node-pair levels focuses on the interactions between arbitrary pairs of nodes. Examples including predicting the distance between all pairs of atoms in molecular graphs.
In Graphi um, um, task at each level starts with data fcaturization, then to the GNN layers (c.g. MPNN Transformer layers), and lastly level-specific MLPs to get each level's embedding. The graph level shares a similar path but uses pooled graph-level features as shown in figure Data featurization such as positional encodings, architecture and metries, and any configurations that are related the levels are easily managed via configuration files.
Multi-task Learning: Graphi um supports multi-task learning where different tasks are learned jointly on the same dataset potentially at different levels. Each level can also support multiple tasks. Each task has its specific task output head (usually constructed as MLP), allowing the model to handle diverse tasks. For instance, at the node level, tasks could range from "node task 1" to "node task N with similar structures available for other levels. Beyond working on tasks within a single level, Graphi um can concurrently optimize tasks across different levels. For example, it can be trained to predict an atom's property (a node-level task) while also determining the type of bond (an edge-level task). Managing the losses and metrics (see section 4. for each task is straightforward, with everything organized within a single configuration file.
Multi-label Learning: Multi-label refers to scenarios where each sample belong more than one class. Unlike traditional classification where each instance is associated with a singular label, in multi-label learning, each molecular instance can be associated with set of labels. For example, particular molecule might simultaneously be therapeutic and toxic, or have both antibacterial and antifungal properties. Graphi im accommodates various tasks, including binary label, multi-labed regression, and multi-class classifications.
E.2 MODELING
Graphi lmm supports straightforward hyper-parameter sweeps, enabling users to fine-tune their models efficiently, improving performance across multiple tasks. Fach task tracked independently using its metrics, ensuring detailed understanding of the model's performance across different tasks. The library supports a variety of GNNs, including 2D, 3D, and graph transformers, providing a seamless user experience and facilitating easy switching between different model architectures.
Expanding Graphium to support more expressive architectures, including transformer models, is critical step in advancing foundation models for drug discovery. The transformer architecture is known for its capacity to capture complex, long-range dependencies and relationships within the data (Rampášek et al. 2022). In molecular graphs, such relationships often contain critical information
25

---

Published as a conference paper at ICLR 2024
Input Features
NILP
Node Feat.
Layers
MLP
Presictions
Rods Taski 2
Edge Feat.
Encoder Manager
Featurized Pya craph
MPNN or Transforner layers
Poo Ling
Graph Tasks
NLP
NLP
Predictioon
Grapn Task
Graph Feat.
Emoco
Conformer Feat.
Edge Tasks
MLP
Prodictions
Edge T254
Positional
MLP
Predictions
Encodings
MLP
Precictsona
Figure 5: Flow chart of multi-level multi-task molecular learning library, Graphium. The Graphiun library offers comprehensive pipeline for molecular learning, consisting of advanced molecular featurization, 2) positional encodings, 3) state-of-the-art GNN and graph transformer layers, 4) multi-level MLPs, and 5) multi-task learning through task specific MLPs the node, edge, and graph levels.
Overview
deep leaningg libery abcaddd graps rppoartatin wanning reaiverrt to
Ctime-ct chetear anchisesuuss.
Pick Ienattration
reebll builhin fexcuttrnion
Premumedd modelts for fast endwaay interncc memate learning
Pysoch
Hore dulavet Grap
interlccs Change
columea atumic hetuutio
Installation
Use AMIDN:
Uretall graptou
Figure 6: Documentation website layout for Graphi um including API references, tutorials, dataset details and more.
about a molecule's properties. Additionally, due to their self-attention mechanism, transformer models provide better interpretability, component much needed in drug discovery. Thus, by supporting more expressive architectures, the library advances the development of foundation models that are not only highly accurate but also interpretable, fostering broader adoption in the scientific community.
To help users understand better the functionalities and APIs in Graphium, constructed an extensive documentation website for Graphi um where API references, tutorials, dataset details and more can be found. This website will be made publicly available and example, include overlook of the website in Figure
E.3 POSITIONAL ENCODINGS
Graphium integrates multi-level positional encodings (PEs) at the node, edge, node-pair, and graph levels into customizable molecular multi-level and multi-task graph learning framework. enables flexible use of various PE combinations, easily managed through compact configuration file. The library introduces novel position transfer functions that facilitate information exchange across positional levels. The most prevalent PEs are the node level and provide nodes with understanding of their positions within a (sub) graph. We further encode relationships between pairs of nodes via edge and nodepair PEs and global graph structure through graph PEs. With modular
26

---

Published as a conference paper at ICLR 2024
design, Graphium supports easy integration of new PEs at all levels, enhancing its ability to handle complex multi-task learning scenarios effectively.
E.4 LABEL NORMALIZATION
Label normalization is well-known technique that improves the regression ability of model by making the labels easier to learn. In a setting with thousands of labels coming from different datasets, the losses associated with each label can be very imbalanced. To counteract that problem, Graphi natively supports multiple normalization schemes, including min-max normalization and 7-score standardization.
E.5 RANKED CLASSIFICATION (HYBRID REGRESSION + CLASSIFICATION LOSS)
For the L1000 datasets, converted the data from regression classification by binning it into categories. To avoid losing the ranking of the classes, we use a hybrid between cross-entropy and MSE loss. Formally, we define this hybrid loss below, with C' denoting the number of classes, X E [0, 1J denoting the prediction, E {0, I,..., C C-1} 11 as the target, and a as the weight of the CE loss,
L(x,0)= = a Σ
1(1 = v) log(t))) + (1 - -a) [ 1=0
()
(x. (1..0) 1) w)²
(1)
1=0
E.6 HANDLING MISSING DATA
An essential component of our multi-tasking paradigm is processing graph data from multiple sources in the same batch. However, graphs in dataset Dj will likely not appear in another dataset D22 resulting in missing labels. In addition, many of the proposed datasets are sparse. Hence, the per-task loss must be computed as the batch consisted exclusively of predictions for which the targets are available.
As a result, Graphium supports different strategies to ignore the missing data or NaNs. These include replacing NaNs with constant number or filtering out NaNs. However, strategies involving filtering do not work when the code compiled, which in turn, necessary for faster training. In these cases, we weigh the missing targets with 0 to compute the desired per-task losses in differentiable way, thereby allowing parallel graph processing from different sources and with missing targets. Specifically, for task with predictions R²T and vector of targets IRT,, we compute
T
C(x,)) T-N N
as the final task loss, where L denotes the per-task loss, T denotes the number of targets and Naan T is the number of nan targets.
E.7 SUPPORT FOR SCALABLE HYPER- PAPPERRRR TUNING WITH μP
To mitigate expensive hyper- pprameee tuning. Graphi um deeply integrates with μP, pre-training paradigm that allows for tuning parameters on a small model, which can be subsequently transferred larger version with out of the box, thus enabling models to scale to billions of parameters efficiently. To the best of our knowledge, this the first work enabling μP for GNNs.
minimal additional tuning required (Yang al.) (022). We support jP
E.8 IPU SUPPORT
To maximize the throughput of training and inference in Graphium, we natively support Graphcore IPUs (Jia et all 2019). The most recent BOW IPUs have 350 TFLOPs of float arithmetic at peak performance and nearly 900MB of high bandwidth on-chip SRAM distributed across 1,472 individual compute cores which can execute programs independently benefit of allowing model weights and activations to be stored directly SRAM but also provides highly flexible, fine-grani gre-grained fabric for compute parallelization. Additionally, this has specific benefits for GNN-like model structures as it allows accelerate sparse communication operations like gathers
all other cores. This has the general
27

---

Published as a conference paper at ICLR 2024
and scatters (Helal al. 2022 and efficiently parallelizes complex combinations of smaller matrix multiplications that can be found in SOTA GNNs (Masters et al. (023). This makes it an ideal hardware platform for the model types targeted by Graphium.
E.9 LIBRARY OPTIMIZATION
Graphi um is optimized for large-scale molecular datasets and models, allowing for fast processing and loading, and is compatible with IPU accelerations while supporting traditional CPU and GPU resources; see below for details.
Dataloading Optimizations are implemented to handle large-scale datasets efficiently, molecules are processed in parallel to apply a set of user-defined featurizers. Additional performance is obtained using a batched- -araeedl- implementation, where each worker processes a batch of molecules, typically on the order of thousands. As the scale of the datasets grows, keeping the data memory becomes a challenge, particularly memory usage can grow linearly with the number of workers. To avoid these issues, in Graphi um, each molecule stored as a separate file on the disk. During data-loading, each worker receives the number of the molecule in the dataset to load, which directly corresponds to the filename to process, massively reducing data serialization and memory usage. Furthermore, storing molecules on disk provides opportunity to skip pre-processing the dataset for every training run.
PE Caching We implement caching of intermediate computations related the PE derivations (e.g.. eigendecomposition of the graph Laplacian and higher-order diffusion operations) that are relevant for the construction of several PEs. This is especially beneficial when large selection of PEs is used in model. Further, we can thereby use positional information across several positional levels (via the position transfer functions proposed in Appendix E.3 with minimal additional computations required.
IPU-specific Optimizations are supported by Graphium um, such multiple device iterations (per- forming multiple training steps on the device, reducing synchronization between the host and the accelerator), buffer prefetch (placing multiple batches on the device ahead of time), IO tiles (dedicat- ing specified number of IPU tiles to data transfer), and loading data asynchronously on the host while the IPU is processing. The Poplar software stack used to run programs on the IPU requires that the shapes of tensors be fixed at compile time, packing strategy also applied to the input data to minimize the padding data required. The library also supports pipelining large models over multiple processors.
Low Precision Training in both mixed precision and true 16 bit floating point precision support in Graphi tm allows reducing training time and memory requirements.
F TOKEN COUNT AND COMPARISON TO GPT-2
Here we compare proposed datasets with datasets used to pre-train foundational models in language, in particular GPT-2, pre-trained on 40GB of text (Radford et al.) 2019). Due to the differences in the learning objective and data domain, direct comparison between the pre-training data of GPT-2 and our datasets is difficult. In addition, the notion of tokens for graph learning models not clearly defined. For example, for graph with n nodes and e edges, one may regard the as the number of tokens. Further, since the edges represent the structural information in a graph, defining the number of tokens as n n might be more meaningful. This notion, however, might still not suffice to fully represent a meaningful token count, as we do not know beforehand the relative importance of node and edge tokens on pre-training graph learning models. Hence, we propose to additionally use the number of prediction targets or labels to compare the token count for two reasons: First, the number of labels has a direct influence on the number of gradient steps we can do per epoch, before the training data repeats. For example, pre-training labels is likely favorable to pre- pre-training on very large graphs with only very few labels. Second, when pre- training transformer, such as GPT-2, on next-word prediction each input token also be regarded as prediction target, making the count of prediction targets an equally meaningful number for comparison.
small graphs with many
28

---

Published as a conference paper at ICLR 2024
Unfortunately, exact token counts for GPT-2 are not available. However, estimate the number of pre-training tokens as follows. GPT-3, the successor model of GPT-2, was pre-trained on 570GB of text, resulting in 400B tokens (Brown et al. 2020b). Crucially, both GPT-2 and GPT-3 use Byte Pair Encodings (Sennrich et al. 2016), meaning that sentences are likely broken up into the same byte-pair encoded tokens. Hence, estimate GiB-tottkenn conversion ratio of
400B tokens
0.7B tokens
570GB
GB
resulting in roughly 28B tokens for GPT-2.
For proposed datasets, present overview the available labels in Table and arrive at total of 13.04B across all mixes. Further, the number atoms and bonds per dataset can be found in T Here, count total of 4.5 atoms and bonds. While these numbers still short of the 28B tokens of GPT-2, we argue that especially our labels can also be seen as more informative for the following reasons.
First, in contrast to the training data of GPT-2, where each label is classification target, our labels comprise both classification and regression targets. In particular, regression targets might provide the model with much more fine-grained information, a factor which should be taken into account.
Second, the joint prediction of multiple labels, potentially even on multiple levels such graph and node-level, might provide additional information during training that is not easily captured with simple count of the labels.
Third, there is evidence in the literature that supervised pre-training outperforms self- supervise methods when the domain adaptation is similar the training domain (Yang al.) 2020)) We expect this to hold in our case since we have gathered a very large diversity of training tasks, increasing the likelihood that least one training label close the desired fine-tuning label. Further, our dataset covers most of the molecules in PubChem, which means that they should cover big part of the chemical space that is usually covered in drug-discover, projects.
G COMPUTATIONAL RESOURCES AND HARDWARE
Hardware Baseline results trained Bow Pod-16 IPU system, made of 16 Bow IPU accelerators, with some additional results obtained using A100 and V100 GPUs.
ToyMIx Computational resources The results presented for the ToYMIX baselines used approxi- mately 142 IPU hours of compute but we estimate that we used approximately 10 times that amount in the development stages for these models.
LarGeMix Computational resources. The results presented for the LARGEMIX used approxi- mately 600 IPU hours (for LARGEMIX. PCQM4M _N4, PCQM4M (000M44_G25) and 150 GPU hours (for PCBA_ 1 1328, L1000 _CAP, L1000_MCF] of compute, with more compute used for development and hyperparameter search.
UltraLARge Computational resources. The results presented for the ULTRALARGE used approximately 500 GPU hours of compute, with more compute used for development and hyperpa- rameter search.
H SCALING LAWS ON LARGEMIX
Table 7: Scaling Law Results for MPNN+ model the proposed LARGEMIX datasets.
AVPR 1
R
model size (hidden size)
1,1000_ VCAP
L1000. MCF7
PCBA_ 1328
PCQM4M GG5
PCOM4M N4
million (64)
4074
.1733 0016
2116 .019
.7737
.7816
million (67)
.4060 :: 2004
.4226 1026
2262 1 .084
.7900
0018
.1899 1: 0006
10 million (128)
.4219 .011
.4475
.2997 1: 1M75
.7974
.505 0017
30 million (225)
.4608 #: .0185
.4872
1002
.3405 + .M22
.7959
0006
.8614 +: 2014
100 million (416)
.4914 :: .013
.5235
0959
.3667 : MM47
.7981
4 0004
.8639 :: 0013
300 million (726)
.5280 +: .004
.5657
1034
.3973
: 0011
.8005
1 0008
.8658
: 0008
billion (1331)
.5285 20055
.5344
.007
.4064
0022
.7995
0007
.8880
0007
29

---

Published as a conference paper at ICLR 2024
Experimental setting. Here aim to investigate the effect of scaling law (number of trainable parameters) on our proposed datasets with the MPNN++ model Masters et al. (2023). We tested the zero-shot scaling and up-scaling using the mup implementation Yang al. (021), meaning that parameters were changed except the hidden dimensions being scaled. To our knowledge, this the first time that mup was used in graph neural networks, its behavior not known. The Graphium library provides access to this valuable tool to the community. making investigation into scaling laws and previously unexplored model sizes for GNNs. To our delight, mup worked very well, showcasing another advantage of Graphium, and we can observe great scaling laws and model improvements up billion parameters.
Results discussion. As shown in Table 7 MPNN++ models with more parameters achieves signif- icantly better results across the range of tasks. The top result balded while the second place is placed with underline. Specifically, on the L1000_ VCAP and 1.1 D00_MCFF) the avPR went up 0.42 + 0.53 and 0.44 0.56 respectively, which a large gain considering that 0.33 would be random classifier. Further, PCBA_ 1328 went up from 0.30 - 0,41, again a large improvement considering the OGB-PCBL leaderboard (a smaller but similar dataset) is at 0.32. further see results going down when the network is scaled down 3M 1M parameters.
Hyperparameters. The hyperparameters for the MPNN+ experiments are specified in Table 8 for reproducibility.
Hyperparameter
Description
GNN depth
16
dropout
0.01
aggregation method
MLP expansion ratio
node embedding combine method GNN normalization
concatenation
batch normalization
any other normalization
layer normalization
gnn output dim
512
virtual node
log gelu
gnn activation
any other activation
relu
Graph output NN graph depth, hidden
2, 512
Graph output NN node depth, hidden task heads depth, hidden batch size
2, 256
2, 128
8192
learning rate
Ir 0.008 with 5 epoch warm up
optimizer epochs
Adam 100
Table 8: Hyperparameters for the MPNN ++ model in the scaling law experiment.
I ADDITIONAL VISUALIZATION ON MOLECULES
In this section, we provide additional visualizations of 24 molecules for L1000_MCRR in Fig-
PCBA_ 1328 1564K in Figure X PCQM4M G25 N4 in Figure 9 QM9 in Figure 10
Tox21 I in Figure II and ZINC 12K in Figure [2
30

---

Published as a conference paper at ICLR 2024
ane
F
gyono
gomo
gono
oar
Figure 7: Molecule visualizations from L1000_VCAP dataset.
mume a A
rr to
546
800
330
anpo
aomg
0400
woroodd
ortapa
Figure 8: Molecule visualizations from 1564K dataset.
31

---

Published as a conference paper at ICLR 2024
amad
Figure 9: Molecule visualizations from PCQM4M_G25 N4 dataset.
Figure 10: Molecule visualizations from QM9 dataset.
32

---

Published as a conference paper at ICLR 2024
x
9o
Figure 11: Molecule visualizations from Tox21 1 dataset.
Qu
fanda
qnob
ane
2nd
Figure 12: Molecule visualizations from ZINC 2K dataset.
33


Reviews:
[{"review_id": "LjHoqbxCuk", "reviewer": "Reviewer_Zdj5", "summary": "This paper present seven novel datasets with different sizes, to facilitate the application and development of molecular foundation model. Multiple baselines are conducted and some experimental results and analysis are provided on several downstream tasks. Furthermore, Graphium graph machine learning library is proposed to offer more convenient implementations of multi-task and multi-level machine learning methods.", "questions": "1. Although the molecular numbers reach the million scale, it appears that the majority of the contribution comes from quantum data. The tasks lack diversity, and there is an imbalance among the three scales of data (toymix, largemix, ultralarge). For instance, the largest ultralarge dataset solely comprises quantum tasks. It would be more comprehensive if it included more biological tasks, such as ADMET properties.\n2. Seems that this paper did not create any new data, instead, they only collected some existing data to organize to different levels or merge different datasets, right? In this case, the novelty and contribution is limited. \n3. From Table 1, we observe that training on multitasks aids in improving the performance of ZINC12k and Tox21, but it doesn't yield the same benefit for Quantum tasks (QM9). What could be the reason behind this phenomenon?\n4. If adding the Quantum tasks enhances the performance of biological tasks, why does the multitask performance for the PCBA_1328 dataset in the largemix dataset at Table 2 not surpass that of the single task setting? The data size may not be the final answer and further analysis may be necessary.\n5. The baselines are too simplistic, as they only include three different types of GCNs. More models involving various types of self-supervised methods, which are essential for a foundation model, should be included in the experiments as baselines for discussion.\n6. If this dataset is aimed at the foundation model, it could potentially achieve the best performance on various tasks(both quantum and bio-related tasks) across different scales. However, this paper only tests various networks on data of different scales dataset, and it lacks related experiments or discussions.\n7. More sophisticated split functions beyond random, such as scaffold, can be explored, as the out-of-distribution (OOD) problem is common in large models and requires careful study.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper compiles a comprehensive benchmark for the evaluation of molecular models. The dataset covers a large variety of molecules and labels at different levels, with particularly comprehensive and complete data for quantum-related tasks.\n2. The Graphium library appears promising as a convenient tool for researchers or developers to perform fine-tuning on various tasks using different backbones.", "weaknesses": "1. The diversity of this benchmark remains a problem. For example, most of the contribution comes from the quantum related property data, which cover a limited part in molecular property prediction.\n2. The paper just collected existing data instead of creating new data, which limits the significance.\n3. Baselines and discussions are not thorough, leading to unconvincing results.\n4. Some empirical settings are not clear or strict, making the comparison or benchmark not solid.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DdjcChUWWu", "reviewer": "Reviewer_jQ36", "summary": "The paper presents seven novel datasets that push the boundaries in both the scale and diversity of supervised labels for molecular learning. The authors also introduce the Graphium graph machine learning library to simplify the process of building and training molecular machine learning models. The paper argues that building effective foundational models for molecular modeling requires supervised training with both quantum mechanical (QM) descriptions and biological environment-dependent data. Overall, this paper aims to explore the possibilities of foundational models in molecular machine learning.", "questions": "- I would be curious about how the models trained on larger datasets perform on smaller datasets (using the same task or different tasks). Also the performance of applying a well-trained model on a small dataset to larger datasets would be interesting to study to understand the few-shot learning ability of those models over these datasets.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents seven novel datasets that are orders of magnitude larger than the current state of the art, covering nearly 100 million molecules and over 3000 sparsely defined tasks, totaling more than 13 billion individual labels, currently the largest of their kind.\n\n- The datasets are designed for the supervised training of foundation models by combining labels representing quantum and biological properties acquired through both simulation and wet lab experimentation. The diversity of labels facilitates efficient transfer learning and enables the construction of foundational models by improving their generalization ability for a wide range of downstream molecular modeling tasks.\n\n- The paper introduces the Graphium graph machine learning library to facilitate efficient training on these extensive datasets. This library simplifies the process of building and training molecular machine learning models.\n\n- The paper provides an anonymized repo for reproducibility commitment.", "weaknesses": "- The evaluated baselines are too simple/out-of-date in some sense. Since this paper is providing a data-driven platform for accelerating the research of molecular foundation models, it would be better to have included more SOTA baselines for comparison, which can help analyze the usefulness and effectiveness of this database better.\n\n- It would be better to have a well-organized website with readable documents that can help understand the details and setups of different datasets/components better.\n\n- More data splitting settings (e.g. mofit-based, system-based, scaffold-based, etc.) should be further studied. Now the OOD generalizability of GNNs across different molecule datasets still seem to be unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "urQMwpeHS3", "reviewer": "Reviewer_qFVJ", "summary": "Having access to a free and open database of chemical compounds with their quantum features and biological activities. The main issue with the existing ones is the shortage of enormous data over various tasks, with many of them lacking experimental measures, or quantum and chemical properties. One know fact is that a little change in these properties can cause a major bioactivity behavior change. To end this, authors has collected a multiple number of datasets over molecules that will be a great help to build foundation models on this literature, e.g., models that are trained over multi-task and multi-level molecular datasets.", "questions": "1. Please provides more samples of each dataset category for the sake of better illustration.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The contribution of this paper has been successfully delivered to the general audience. It will open doors to having more robust and representative models with the help of pre-training. \n2. Advantages over non supervised methods for pre-training the model, such as contrastive learning, is important because of the \"activity cliffs\" phenomenon. \n3. Many intricacies are considered into this, such as having both quantum and bioactivity features which are important for the future research. \n4. Graphium library provides a reliable and tidy tool that gives a better pace to doing research.", "weaknesses": "1. It's not clear what are exact labels collected for each sample and how they can be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Hzy2uRPXHN", "reviewer": "Reviewer_BQyU", "summary": "The authors have developed an extensive data repository and designed a graph learning library, “graphium” that supports the training of foundational models for molecular modeling. Authors have combined molecular datasets from various sources with labels for tasks targeting graph and node-level properties. The evaluation shows that training graph-based models on large amounts of quantum mechanical data improves the downstream performance of the model in resource-starved contexts, for example, modeling biological molecules.", "questions": "+ Authors mention that due to resource constraints, they only use a small fraction of the data to train the ULTRALARGE models; can they discuss how would the results change if the entire dataset is used? As mentioned by the authors, there is a skew towards the QM tasks. \nWould the model overfit those tasks when trained with larger datasets and show degradation in performance for biological tasks? \n\n+ It's unclear if the paper shows that the model can generalize to unseen tasks in the downstream setting. \n\n+ Minor point: I am curious to see how the newer equivariant GNNs [3, 4] (specifically in the molecular modeling space) interact with having more data available to them.\n\n\nReferences\n\n[3] E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials https://arxiv.org/pdf/2101.03164.pdf\n[4] Equivariant Graph Attention Networks for Molecular Property Prediction https://arxiv.org/pdf/2202.09891.pdf", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "+ The paper provides an extensive description of all the datasets that have been used and how they relate to each other in terms of what tasks they model and the labels provided. Particularly, the paper provides an in-depth explanation of (in the supplementary resources) how chemical and structural properties might relate to biological function, which can explain how properties learned in a quantum mechanical setting might help in a biological context. \n\n+ Authors have made their work publicly available, including the datasets and the graph learning library, and providing multi-level, multi-label, and multi-task learning mechanisms. This form of modeling, as shown, is crucial in developing foundational models that generalize well to tasks with limited amounts of data available.  \n\n+ The authors discuss the positional encodings and distill the literature on relative and global positioning into their “graphium” data modeling pipeline.", "weaknesses": "+ The work is an important contribution to the ML and molecular learning communities. However, it restricts itself to data curation and performance results for different GNN-based models. The paper is missing crucial insights that might help inform the researchers interested in developing foundational models. For example, does the new dataset help improve other existing methods? Why the multi-task setting does not work well for all the datasets, and what are the data properties resulting in different performances for different datasets? How does this work translate into applications for the real-world setting? \n\n+ The paper does not show a comparison against any other existing pre-trained molecular models. Moreover, it would also be helpful to highlight the performance against supervised SOTAs.\n\n+ Although the paper mentions transformer-based models, such as Graphormer or GPS [1,2], in the supplementary section, it does not provide any results for them in the evaluation section. Showing results for them may be an essential consideration because they would show potentially better scaling to data than standard message passing, which tends to suffer from the problem of over-smoothing when the number of layers goes beyond three. Although the paper mentions this limitation and plans to do the comparison in the future, it's crucial to define how transformer-style models scale with the data that has been curated. \n\n+ Continuing on the previous point, it would be nice to see a scaling curve that shows how increasing model complexity improves performance as we provide more data to it.\n\n\nReferences:\n[1] Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets arXiv:2203.04810\n[2] Recipe for a General, Powerful, Scalable Graph Transformer arXiv:2205.12454", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "shRPd8uu46", "reviewer": "Reviewer_43gK", "summary": "This manusript proposed a curated dataset containing three subsets with different sizes. The whole datasets contains approximately 100 million molecules with more than 13 billion labels combining graph-level and node-level ones. Each molecules is accompanied with quantum mechancial properties calculated using either DFT or semi-emperical methods. A Python library called Graphium is built to facilitate the access and utilization of the proposed dataset for machine learning applications. Some tools such as handling missing data and label normalization in Graphium are designed to reduce the friction for users. Lastly, the authors provides benchmark results on the proposed dataset using common GNN models.", "questions": "1. For the PM6_83M, the authors mention that the QM properties of molecules are calculated using semi-empirical methods as reported in the original paper of PM6 paper. However, in secion D.4, the authors state that \"Using the 3D conformation optimized by the DFT, we further computed the following 3D descriptors as labels\". Can the authors confirm if they run DFT on top of the PM6 dataset or not? If they do run DFT, can they briefly discuss the DFT method they use here?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed dataset is large and inclusive. In each one of the subset (ToyMix, LargeMix, and UltraLarge), at least one of the dataset contains 3D conformer of molecules and QM properties. As the 3D GNN models (e.g. Equivariant GNNs) are becoming increasingly important in the molecular machine learning community, 3D molecular conformer and corresponding labels are valuable.\n\n2. The Graphium library can be handy for researchers who wants to join the molecular machine learning community but are halted by lack of experience with data processing and model training. With all datasets curated in this work included in the Graphium library, users can easily focus on innovating model architectures. \n\n3. Many node-level labels are included in this dataset. Most conventional tasks in molecular machine learning are graph-level tasks such as molecular property prediction because of the scarcity of node-level labels.", "weaknesses": "1. One minor weakness is that the potential noise in the bioassay data. For the PCBA1328 dataset of LargeMix, the authors curated the data from PubChem by collecting molecules and their properties with regard to different bioassays. Due to the inherent noise associated with the data generation process with bioassay, model training (especially multitask training) can be difficult.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)