On the Scalability of GNNs for Molecular Graphs

We thank the reviewer for providing detailed feedback on the paper which is of utmost value to our work. Below we address your concerns point-by-point. We also kindly invite the reviewer to refer to the general rebuttal where we share further information and a summary of the feedback we received.

The pre-training strategy [is] supervised [...], overlooking a range of existing self-supervised approaches [...]

Thank you for your valuable feedback. We agree that our work utilizes a lesser explored strategy of supervised pretraining. Unsupervised molecular models are built on years of cutting-edge pretraining research [2,3,4] with vast empirical promises [6]. In fact, our work, to our best knowledge, is the first to successfully scale supervised pretraining for GNNs to the multi-billion parameter regime, setting new downstream task performance standards for the multiple benchmarks considered in our submission.

In Table 1 and 2 of the rebuttal doc, we compare our approach to self-supervised strategies and find that our method maintains the clear upper hand. The unsupervised MolE [10] model variants clearly underperform the standard MolE model that leverages both unsupervised and supervised pretraining (Table 2), which is in turn outperformed by the listed MolGPS variants by large margins. Even our smaller MPNN++ models of comparable size to MolE (~100M parameters) both outperform the self-supervised variant.

In Table 1, we compare to the unsupervised GraphMVP [11] on MoleculeNet (only on tasks that were not part of our pretraining). Our MolGPS from the original submission outperforms [11] on 3/4 tasks, while the model variants that we recently pretrained with additional Phenomics data outperform across all tasks. We acknowledge the results reported by Uni-Mol [12] slightly surpass our score, but note that [12] uses a different data splitting approach compared to GraphMVP, which can have a significant impact in low-data regiemes. We have checked the paper and code, but we were unable to find the recipe so far.

Lastly, the mentioned reference [13] focuses on 3D molecular modeling, which is outside the scope of the present work (as discussed below).

The pretraining data [...] of only 5 million molecules.

We would like to point out that due to the supervised pretraining, the scale of the dataset is not directly comparable. We also note that previous works in GNN literature scale to less than 5M graphs [14] and show interesting scaling trends.

In our case, it is important to characterize the diversity not only by the number of molecules, but instead paired with the number of labels per molecule. We recall our label scaling experiments that show the impact of reducing the data diversity by removing labels. The performance gains from incorporating Phenomics data further support this hypothesis (Figure 2 of rebuttal).

The scaling of molecular data appears to perform poorly [for] the downstream [...].

We indeed found no consistent trend on downstream task performance for increased number of molecules in the pretraining dataset (Figure 2(a) and Appendix E.3 of our submission). That said, no molecular scale (except for the lowest 12.5%) had an outstanding negative effect on downstream tasks. We therefore hypothesize that our approach is robust to a smaller number of molecules during pretraining when the highly important label diversity per molecule is maintained (as discussed above).

[I]t is also important to evaluate quantum molecular properties such as QM9 and MD17.

Thank you for the suggestion. However, datasets such as QM9 and MD17 require the model to reason over graph geometry in 3D coordinate spaces. Evaluating on such tasks would be an uneven evaluation of a 2D model such as MolGPS being used to learn 3D downstream tasks.

We would like to also highlight the strong performance of MolGPS on the complementary gene inhibition tasks (e.g., the pkis2-* task series of Polaris benchmark) and the significant improvements of our most recent model variants that have been pretrained with additional Phenomics data (Figure 1 in rebuttal pdf).

Furthermore, ADME(T) tasks capture the efficacy and toxicity of bioassays which play a crucial role in screening and evaluation of likely drug candidates. These tasks continue to remain the pinnacle of industrial pharmaceutical evaluation [15,16]. We thus prioritize our empirical evaluation by keeping their real-world applications and biological importance in mind.

[2]. Li et al, A knowledge-guided pre-training framework for improving molecular representation learning, Nature 2023

[3]. Xia et al, Pre-training Graph Neural Networks for Molecular Representations: Retrospect and Prospect, ICML AI4Science Workshop 2022

[4]. Li et al, KPGT: Knowledge-Guided Pre-training of Graph Transformer for Molecular Property Prediction, ACM SIGKDD 2022

[6]. Lu et al, Learning to Pre-train Graph Neural Networks, AAAI 2021

[7]. Sun et al, Does GNN Pretraining Help Molecular Representation?, NeurIPS 2022

[8]. Sun et al, MoCL: Data-driven Molecular Fingerprint via Knowledge-aware Contrastive Learning from Molecular Graph, ACM SIGKDD 2021

[10] Méndez-Lucio et al, MolE: a molecular foundation model for drug discovery, arxiv 2022

[11] Liu et al, Pre-training molecular graph representation with 3d geometry. arXiv preprint arxiv 2021

[12] Zhou et al, Uni-mol: A universal 3d molecular representation learning framework, ChemRxiv 2023

[13] Feng et al, Fractional denoising for 3d molecular pre-training, ICML 2023

[14]. Chen et al, Uncovering neural scaling laws in molecular representation learning, NeurIPS, 2024

[15]. Shi et al, Fine-tuning BERT for automatic ADME semantic labeling in FDA drug labeling to enhance product-specific guidance assessment, Journal of biomedical informatics, 2023

[16]. Walter et al, Multi-task ADME/PK Prediction at Industrial Scale: Leveraging Large and Diverse Experimental Datasets, Molecular Informatics, 2024

We hope our two comments above addressed the reviewer’s concerns. Please let us know if there are any additional clarifications that we can provide that could increase your support for the acceptance of the paper.

We thank the reviewer for providing detailed feedback on the paper which is of utmost value to our work. Below we address your concerns point-by-point. We also kindly invite the reviewer to refer to the general rebuttal where we share further information and a summary of the feedback we received.

(W1) Figure 2 is too small and difficult to read

We thank the reviewer for the comment and agree that Figure 2 can be better represented when split up into multiple larger figures. We will apply those changes in a revised version of our paper.

(W2 & Q4) Fine tuning vs probing: How does fine-tuning compare to probing?

We thank the reviewer for this helpful comment. We will further clarify this point when revising the paper.

Section 4.2 of the submission studies finetuning and probing side-by-side, establishing both as effective strategies for tackling downstream tasks, without conducting a direct comparison between them.

However, when deriving the foundation model MolGPS in Section 4.3, we find probing to be the overall stronger approach. The major advantage of probing is the ability to leverage multi-level information from the pretrained GNN. We recall that our pretraining is based on a supervised multi-task learning approach. As a result, different task heads capture task-specific information, while earlier layers that feed into the task heads carry more general information. When we combine fingerprints from various layers, we can think of aggregating knowledge from several “experts”. Our multi-fingerprint probing suggests this knowledge is additive as it clearly outperforms probing of any single fingerprint. Our foundation model MolGPS doubles down on this idea, combining fingerprints from multiple layers and various pretrained models. For finetuning, there is no straightforward way of taking advantage of this multi-level information, making probing the preferred approach.

(Q1) How can you train a single model on multiple datasets if these datasets might have different features? Are all datasets pre-processed to have the same features?

All datasets are indeed pre-processed to have the same features. Each molecule is initially represented by a SMILES string [1], that our pipeline features into a molecular graph (i.e., atoms and bonds as nodes and edges, respectively). The (node) features are also agnostic to the data source, e.g., using atom type, row/col in periodic table, etc.. We will add more context on the pre-processing to the appendix.

(Q2) Figure 2 - probing on Polaris / TDC: it seems that most models do not profit from an increase in the number of molecules. Does this not contradict the idea of foundation models?

We indeed found no consistent trend on downstream task performance for increased number of molecules in the pretraining dataset (Figure 2(a) and Appendix E.3 of our submission). That said, no molecular scale (except for the lowest 12.5%) had an outstanding negative effect on downstream tasks. We therefore hypothesize that our approach is robust to a smaller number of molecules during pretraining when the highly important label diversity per molecule is maintained that can be observed in our label scaling study (e.g., in Figure 2(a) and Appendix E.4 of our submission).

We would like to point out the importance of characterizing the diversity of the data not only by the number of molecules, but also paired with the number of pretraining labels/tasks per molecule. PCBA_1328 considers more than 1k different labels per molecule (albeit with high sparsity) and PCQM4M comes with 25 graph-level tasks and 4 node-level tasks (that are learned for each node of the ~4M molecules). This perspective is reinforced by our label scaling experiments that show the impact of reducing the data diversity by removing labels (e.g., Figure 2(a) of submission). Performance gains from recently incorporating Phenomics data into the pretraining data mix further support this claim, adding ~500k of molecules with a highly informative set of ~6k labels per graph (Figure 2 of rebuttal).

(Q3) In a similar vein, you notice that you can get better results when pretraining without the L1000 dataset. Does this imply that for graphs it is more important to have the right kind of data compared to having a lot of data (as for example LLMs or diffusion models)?

This is an accurate observation, which relates back to the importance of the number of molecules and the number (and quality of the labels). L1000 is exceptionally small (only 20k molecules) but features almost 1k labels for each molecule. However, in contrast to, e.g., PCBA_1328 (1328 binary classification tasks), it is a regression task that suffers from low signal-to-noise ratio. Despite trying to stabilize learning by transforming it into a classification task via binning of regression targets, the overall impact on downstream task performance was negative. Notably, the small amount of molecules may lead to the absence of positive samples in most batches, as the batches are dominated by the larger PCBA_1328 and PCQM4M molecules.

As pointed out by the reviewer, it is the quality (and label diversity) of the data that is highly relevant. The perfect example for this is the massively improved downstream task performance we report in the rebuttal doc after recently adding Phenomics data (with ~500k molecules and ~6k highly informative labels) to our previously use pretraining data mix (excluding L1000). We kindly refer the reviewer to the general rebuttal doc for more details.

Kindly let us know if our response above addressed your concerns. We will be happy to further discuss and answer any questions you may have.

[1] Weininger D, SMILES, a chemical language and information system. 1. Introduction to methodology and encoding rules. Journal of chemical information and computer sciences. 1988

We thank the reviewer for their valuable feedback on the paper which is of utmost value to our work. Below we address your concerns point-by-point. We also kindly invite the reviewer to refer to the general rebuttal where we share further information and a summary of the feedback we received.

Another important aspect of scaling the parameters is regularization [1]. The authors did not explicitly discuss / validate the regularization methods they are using.

Our model indeed uses regularization techniques, e.g., dropout in each GNN layer. We will include a detailed discussion of the used methods and their parameterization in the appendix of the revised paper. The need of more sophisticated regularization techniques is largely avoided through the use of rich pretraining data that prevents the models from overfitting (except for the MPNN++ model in a few instances).

The authors did not control the compute used to train the models. It would nice to see some efficiency comparison as well.

Thank you very much for the suggestion. Model performance tables with the training compute time will be added to the appendix. Training time of a single model varies from 24 GPU hours for the smallest models to ~1120 GPU hours for the largest models. In total ~13400 GPU hours were used for computation in this paper.

The work of [1] pointed out by the reviewer presents a compelling analysis of the used computational resources compared between finetuning to a downstream task and training a model of the same size from scratch.

We note that such a study may be difficult to replicate in our case due to the high parametric scale of the GNNs in contrast to the low-data regimes for the downstream tasks with only a few 100's or 1,000's of molecules per task and almost exclusively a single target label, which would cause overfitting when trained from scratch.

How well do these insights generalize to other types of graph-structured data beyond molecular graphs?

We thank the reviewer for their question on potential applications outside the biochemical domain. It is highly likely that models similar to MolGPS could be pretrained for other applications and domains given adequate pretaining data. However, the present work would not be suited for downstream tasks outside the biochemical domain. To obtain a domain-agnostic graph-foundation model, unsupervised approaches that primarily learn from the graph structure may be a more natural choice.

That said, the proposed foundation model MolGPS is specifically targeted towards biochemistry and drug discovery in particular, with the pretraining data chosen accordingly to learn domain-specific information. For comparison, we provide some comparison to unsupervised pretraining approaches in the rebuttal doc (Tables 1 & 2) that indicate our supervised pretraining strategy is to be preferred for a “molecular” graph foundation model.

Kindly let us know if our response above addressed your concerns. We will be happy to further discuss and answer any questions you may have.

[1] Steiner et, How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers, arxiv 2021

We thank the reviewer for providing detailed feedback on the paper which is of utmost value to our work. Below we address your concerns point-by-point. We also kindly invite the reviewer to refer to the general rebuttal where we share further information and a summary of the feedback we received.

Many scaling results are unsurprising, such as the consistent improvement with increased depth and width during pretraining. However, the thorough experiments compensate for this predictability.

Thank you for appreciating the quality of our empirical analysis and finding our findings interesting.

We agree that, especially given the unprecedented benefits of scale in other domains such as natural language, some results of this work may seem unsurprising at first. In the context of GNNs however, our work explores largely uncharted territory in terms of parametric scale [1,2,3] and the possible implications on biochemistry and drug discovery in particular. We recall the finetuning and probing performance on downstream tasks, establishing new state-of-the-art on 12/22 tasks on highly-competitive TDC benchmark and all but two tasks on Polaris and MoleculeNet benchmarks with our proposed foundation model (MolGPS). In our rebuttal we have expanded our empirical study, adding model variants of MolGPS that were scaled to the 3B parameter regime andpretrained on an improved data mix that further enhances our performance across the downstream task benchmarks. We kindly refer to the general rebuttal for more details.

I am curious if these trends will also occur for 3D graph datasets like QM9 or others. Conducting a few additional experiments on these datasets could strengthen this paper, though this isn't necessarily required in this rebuttal.

Thank you for the suggestion. We note that our proposed MolGPS model solely relies on learning from 2D graphs, while tasks in QM9 depend on molecules with 3D positions. This is one of the reasons why we limit our analysis of downstream tasks to 2D molecular tasks.

It is important to highlight the importance of publicly available labeled data for our supervised pretraining setup. Our label scaling experiments (e.g., Figure 2a of submission) highlight the importance of many diverse labels for each molecule. Current databases for 3D molecules are still fairly small and would hardly allow for the required data diversity present in our 2D pretraining data collection. However, we note that scaling design decisions uncovered in our empirical analysis may prove useful to construct and train foundational 3D GNNs in the future, when databases reach adequate scales. These models would further require the use of sophisticated training strategies and aggregation functions, a future direction which we explicitly mention in our limitations section.

Kindly let us know if our response above addressed your concerns. We will be happy to further discuss and answer any questions you may have.

[1]. Sun et al, Does GNN Pretraining Help Molecular Representation?, NeurIPS 2022.

[2]. Sun et al, MoCL: Data-driven Molecular Fingerprint via Knowledge-aware Contrastive Learning from Molecular Graph, ACM SIGKDD conference on knowledge discovery & data mining 2021.

[3]. Liu et al., Neural Scaling Laws on Graphs, arxiv 2024.