Paper: Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
We thank the reviewer for pointing out this issue. Indeed, due to the nature of the data, there will be a base level of noise in the bioassay data. In this work, we curate the bioassays from existing data and utilize the reported flags from original authors. To denoise the data, the active and inactive labels from the original data are treated as 1 and 0 respectively while everything else is considered as NaN thus ignored rather than relying on the specific measured values which are more subject to experimental noise. Lastly, performing classification tasks on bioassay data helps reduce the noise by simplifying the labels to be binary though some effect of mislabeling might still be present.

Reviewer comment (Reproducibility):
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?

Author response:
We thank you for reporting this error, we will correct it. The PCQM4M dataset has conformations minimized by DFT methods, with the labels also being computed using the same DFT. However, the PM6_83M dataset minimizes the conformations with semi-empirical PM6 method, and uses the same semi-empirical method to compute the labels.

Reviewer comment (Evaluation):
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?

Author response:
We thank you for these very important points. First, we want to reiterate that as part of the contributions of this work we do not just provide the curated datasets, but also Graphium, a library to enable large-scale training for molecules with various hardware. The Graphium library supports many SOTA models and positional encodings, mup for zero-shot scaling, and TDC for fine-tuning. The library being very modular, it allows to easily swap architectures, datasets, positional encodings, metrics, and benchmarks specifically to help study the very questions you are proposing. 

However, most of these questions are left for future work, as our paper proposes the first step to bring us there.

**Why multi-task does not always work well**

Regarding the question on why multi-task doesn’t work well on all datasets, we believe the main cause is underfitting, as a 10M parameter model is used to train on >3000 tasks jointly. This is supported by the reported results below, on single task, the test set is between 11% and 30% **worse** than train, but on multi-task, the test set is between 5% and 9% **better**. Again, looking at the scaling results, we can see that the performance on the PCBA_1328 goes up very significantly despite the multi-tasking.

These points were clarified in the paper to better explain the disparity between the single-task and multi-task for PCBA_1328.

|  | Train BCE Single task | Train BCE Multi-Task | Test BCE Single task | Test BCE Multi-Task |
|--|--|--|--|--|
| GCN | 0.0284 ± 0.0010 | 0.0382 ± 0.0005 | 0.0316 ± 0.0000 | 0.0349 ± 0.0002 |
| GIN | 0.0249 ± 0.0017 | 0.0359 ± 0.0011 | 0.0324 ± 0.0000 | 0.0342 ± 0.0001 |
| GINE | 0.0258 ± 0.0017 | 0.0361 ± 0.0008 | 0.0320 ± 0.0001 | 0.0341 ± 0.0001 |

Reviewer comment (Experiments):
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.

Author response:
This is an important suggestion, but the aim of our current work is not to provide pre-trained models, but rather the datasets and library required to build foundational GNNs and benchmark them on TDC. Hence, our current baselines are not meant to be used for fine-tuning, and cannot be compared directly to self-supervised models.

Reviewer comment (Experiments):
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.

Author response:
Referring to the general comment, we have shown some great scaling trends of up to 1B parameters with continuously improving results using MPNN++, the backbone behind the GPS++ winner of the OGB-LSC competition. We also aim to provide some baseline Transformers at the 10M parameter scale. However, in-depth comparison of the scaling laws of MPNN vs Transformer vs Hybrid is left for future work.

Further, there are multiple literature showing that positional and structural encodings are more important than the GNN models itself, such as GPS[1], GPSE[2], LSPE[3], etc., and all models tested in the baseline use eigenvectors and random-walk encodings that are again behind the OGB-LSC competition win.

[1] Rampášek, Ladislav, et al. "Recipe for a general, powerful, scalable graph transformer." Advances in Neural Information Processing Systems 35 (2022): 14501-14515.

[2] Liu, Renming, et al. "Graph Positional and Structural Encoder." arXiv preprint arXiv:2307.07107 (2023).

[3] Dwivedi, Vijay Prakash, et al. "Graph neural networks with learnable structural and positional representations." arXiv preprint arXiv:2110.07875 (2021).

Reviewer comment (Experiments):
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.

Author response:
We thank you for this suggestion, but we leave scaling laws analysis for the data for future work. In the general comment, we provided scaling laws for the model sizes.

Reviewer comment (Experiments):
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. Would the model overfit those tasks when trained with larger datasets and show degradation in performance for biological tasks?

Author response:
This is a great point that we have carefully thought about in the design of Graphium. Indeed, regarding the question of imbalance and skew, Graphium currently supports the parameter `epoch_sampling_fraction` for each dataset. This parameter allows to sub-sample only a fraction of a dataset at any given epoch and can be used to re-balance the amount of molecules seen for each task. Therefore, we believe we can tune this parameter to avoid degradation on biological tasks. 

For example, selecting it to be `=0.05` would yield to about the same number of QM datapoints for both the UltraLarge and LargeMix datasets at each epoch, but would also yield to better generalization since it will take 20 epochs to see the entire dataset.

In the interests of robustness and confidence of results provided the core experiments were priorities, and some experiments were not conducted due to finite compute resources. However, we agree this is an interesting future direction to explore.

Reviewer comment (Experiments):
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.

Author response:
The Graphium library is meant for 2D GNNs, not 3D-GNNs or E-GNNs, so direct comparison cannot be made. Here, the QM dataset only comprises a single lowest-energy conformation, which is used to complement the information of the biological tasks. This is not suited for MD, which requires learning energies and forces outside of local-minima, and which requires the full 3D structure as input.

Reviewer comment (Experiments):
It's not clear what are exact labels collected for each sample and how they can be useful.

Author response:
Description of the labels are in `Appendix D DEEPER DIVE INTO THE LARGEMIX AND ULTRALARGE DATASETS`. For the PCBA_1328, it is impossible to provide a description of all 1328 labels, but the pubchem assay ID is still available in the column header of the CSV file so they are all traceable.

If you have any specific questions on any specific dataset or label that Appendix D does not answer, please let us know.

Reviewer comment (Presentation):
Please provides more samples of each dataset category for the sake of better illustration.

Author response:
In Figure 2 of Appendix C, we provide 6 molecules per dataset. But to answer your request, we will re-do the figure to provide 24 molecules per dataset and span it across 2 pages.

Reviewer comment (Experiments):
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.

Author response:
We would like to point the reviewer to the general comment. We have added a baseline for MPNN++ (the backbone of the GPS++ and winner of the OBG-LSC competition), and aim to provide results for graph Transformer. Further, there are multiple literature showing that positional and structural encodings are more important than the GNN models itself, such as GPS[1], GPSE[2], LSPE[3], etc., and all models tested in the baseline use eigenvectors and random-walk encodings that are again behind the OGB-LSC competition win.

[1] Rampášek, Ladislav, et al. "Recipe for a general, powerful, scalable graph transformer." Advances in Neural Information Processing Systems 35 (2022): 14501-14515.

[2] Liu, Renming, et al. "Graph Positional and Structural Encoder." arXiv preprint arXiv:2307.07107 (2023).

[3] Dwivedi, Vijay Prakash, et al. "Graph neural networks with learnable structural and positional representations." arXiv preprint arXiv:2110.07875 (2021).

Reviewer comment (Reproducibility):
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.

Author response:
We thank the reviewer for this proposition, and would like to mention that we do have a website! The website contains different tabs (Overview, Baseline, API, Tutorials, Design, Datasets, Pretrained Models, CLI, Contribute, License).

For anonymity reasons, we cannot share the website but invite you to build it locally:

- Install mkdocs package `pip install mkdocs`
- Go to the main graphium directory
- Run it on a local server `mkdocs serve`
- CTRL+Click on the link given after the compilation

Reviewer comment (Experiments):
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.

Author response:
We thank the reviewer for this proposition, but would like to reiterate that the objective of the dataset is to pre-train a large model, not to test the OOD. However, we do support direct finetuning on TDC to test the generalizability of the model, but this is left for future work. A tutorial on how to finetune on TDC is available in the file graphium/notebooks/finetuning-on-tdc-admet-benchmark.ipynb.

Reviewer comment (Experiments):
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.

Author response:
We respectfully disagree with the reviewer. We believe our **curated** datasets are highly diverse in the non-quantum tasks as well. For instance, the PCBA_1328 dataset combines 1328 assays covering a vast task space and the L1000 data combines 998 gene perturbations across 2 cell lines. Hence, there’s a total of 3324 bio-assays in the LargeMix dataset, which we curated ourselves into ML ready format fitting for GNNs.  To the best of our knowledge, these alone are more diverse than any public 2D molecular learning datasets in the literature.

It is also worth noting that due to the time intensive nature of results derived from experimental, rather than computational sources, such datasets are significantly smaller than their computational counterparts. While the reviewer is absolutely correct that if more experimental / biological data were available it would be preferable, we argue that the combination of large quantum datasets augmented with a smaller portion of high quality experimental data allows a model to learn suitably general embeddings, and believe the results presented show this.

Reviewer comment (Novelty):
The paper just collected existing data instead of creating new data, which limits the significance.

Author response:
First, we note that the datasets are only one of the contributions, as we also propose an open-source *Graphium* library that contains the SOTA graph network models for molecules and positional/structural encodings.

Second, in the field of molecular learning, collecting novel data, especially for non-quantum tasks is extremely expensive and often involves costly real-world experiments. In addition, available data is often not in a ML ready format and requires extensive data processing with domain experts to prepare them for training in graph neural networks. In this work, we curated 3 datasets for the LargeMix and 1 dataset for the Ultra-large, with labeled features as a starting point for building foundational models on molecules. The datasets cover nearly 100 million molecules and over 3000 sparsely defined tasks totaling more than 13 billion individual labels. There is significant value in the curation of these data in simple tabular formats.

Reviewer comment (Experiments):
Baselines and discussions are not thorough, leading to unconvincing results.

Author response:
The focus of this work is to construct the largest public collection of molecular data and the Graphium library which supports graph ML on the proposed dataset. Our objective is not for building a benchmark to compare many GNN architectures.
The baselines provided are intended to demonstrate both the soundness of the datasets and multi-task learning in a machine learning context, and provide a reference point for other researchers developing new models for expected performance of simple and reproducible models. 

Although more thorough experiments are left for future work both internally and to the community, we aim to add results for the Graph Transformers. We further added some scaling law results using the MPNN++ architecture, the backbone of the GPS++ model that won the OGB-LSC competition, and we observed an impressive scaling trend up to 1B parameters.

Reviewer comment (Experiments):
Some empirical settings are not clear or strict, making the comparison or benchmark not solid.

Author response:
We provide details on our experimental settings for TOYMIX, LARGEMIX and ULTRALARGE in detail in Section 4.1, 4.2 and 4.3 respectively. Appendix E also provides details on the graphium library for multi-level, multi-task and multi-label learning, the modeling and more. We would be glad to further clarify any specific concerns for the reviewer.

Reviewer comment (Experiments):
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.

Author response:
We want to point out that our contributed datasets contain a total over 3000 tasks thus are highly diverse in nature. When examining the balance between biological tasks and quantum tasks, one must first consider the data availability. Note that ADMET properties are generally only found in very small amounts as it requires lab experiments. The well-known TDC benchmark[1] contains a number of ADMET properties however each property only has 500-6000 molecules and are notorious for their curation issues, see [“Getting Real with Molecular Property Prediction](https://tinyurl.com/mr2mvavt) and [“We Need Better Benchmarks for ML in DD”](https://tinyurl.com/cdhht757).

Therefore, we omitted the ADMET at the moment due to the issues with dataset curation and sizes. We also reiterate that we are proposing datasets for pre-training at a large scale, not a task-specific benchmark.

[1] Huang, Kexin, et al. "Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development." 2021.

Reviewer comment (Novelty):
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.

Author response:
We refer the reviewer to the detailed response we provided for W2.

Reviewer comment (Evaluation):
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?

Author response:
We thank the reviewer for raising this point. First, note that performance benefits on the ZINC12k and Tox21 are **more interesting** than QM9. This is because, in practice, biological tasks have significantly less training data than quantum tasks, and the goal of building a foundational model with a 2D GNN is to improve the performance of low-data biological tasks through the joint pre-training with quantum tasks which we observe in this case. 

Second, the models use the same amount of parameters to train on both multitask and the sole QM9 task thus the parameters are shared across other tasks for the multitask model. Third, QM9 originally has a much higher volume of data, so it benefits less from the other tasks.

Reviewer comment (Evaluation):
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.

Author response:
We thank the reviewer for this observation and will clarify it in the paper. We believe the main cause is underfitting, as a 10M parameter model is used to train on >3000 tasks jointly. Thus, we argue that longer training (capped at 100 epochs) and more parameters (capped at 10M) will benefit the multi-task model more.

This hypothesis is supported by the reported results below, on single task, the test set is between 11% and 30% **worse** than train, but on multi-task, the test set is between 5% and 9% **better**. Again, looking at the scaling results, we can see that the performance on the PCBA_1328 goes up very significantly despite the multi-tasking.

|  | Train BCE Single task | Train BCE Multi-Task | Test BCE Single task | Test BCE Multi-Task |
|--|--|--|--|--|
| GCN  | .0284 ± 0.0010  | .0382 ± 0.0005 | .0316 ± 0.0000 | .0349 ± 0.0002 |
| GIN  | .0249 ± 0.0017  | .0359 ± 0.0011 | .0324 ± 0.0000 | .0342 ± 0.0001 |
| GINE | .0258 ± 0.0017 | .0361 ± 0.0008 | .0320 ± 0.0001 | .0341 ± 0.0001 |

Reviewer comment (Experiments):
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.

Author response:
We would like to refer the reviewer to the general comment, where we added baselines for MPNN++, the backbone model behind the OGB-LSC challenge winner, and scaling laws. 
And additionally clarify for the reviewer that the aim is to provide baselines for the datasets, not to optimize performance on this new dataset collection.  .

Reviewer comment (Experiments):
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.

Author response:
We would like to reiterate that we are not proposing a benchmark paper, but rather datasets to pre-train a model, and pre-training is known to work better with random split as it allows the model to better see the desired chemical space. However, the Graphium library does include an easy access to the TDC API where OOD generalization can be tested using the various splits proposed. A tutorial on how to finetune on TDC is available in the file graphium/notebooks/finetuning-on-tdc-admet-benchmark.ipynb.