On the Scalability of GNNs for Molecular Graphs

Scaling deep learning models has been at the heart of recent revolutions in language modelling and image generation. Practitioners have observed a strong relationship between model size, dataset size, and performance. However, structure-based architectures such as Graph Neural Networks (GNNs) are yet to show the benefits of scale mainly due to the lower efficiency of sparse operations, large data requirements, and lack of clarity about the effectiveness of various architectures. We address this drawback of GNNs by studying their scaling behavior. Specifically, we analyze message-passing networks, graph Transformers, and hybrid architectures on the largest public collection of 2D molecular graphs. For the first time, we observe that GNNs benefit tremendously from the increasing scale of depth, width, number of molecules, number of labels, and the diversity in the pretraining datasets. We further demonstrate strong finetuning scaling behavior on 38 highly competitive downstream tasks, outclassing previous large models. This gives rise to MolGPS, a new graph foundation model that allows to navigate the chemical space, outperforming the previous state-of-the-arts on 26 out the 38 downstream tasks. We hope that our work paves the way for an era where foundational GNNs drive pharmaceutical drug discovery.

Paper

References (93)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer Lxfs3/10 · confidence 5/52024-06-24

Summary

This paper investigates the scaling of GNNS across various settings including width, depth, number of molecules, number of labels, and diversity in pretraining across 12 datasets. Different conclusions were drawn from experiments conducted during both pretraining and fine-tuning stages. Finally, the authors introduce a foundational model named MolGPS which integrates various findings.

Strengths

1. The experiments are conducted extensively, taking into account various variables, resulting in some interesting findings. 2. The problem addressed in this paper—the scaling law problem in the molecular domain—remains unresolved but is crucial for advancing molecular representation learning.

Weaknesses

1. The pre-training strategy employed in this paper is a supervised method, overlooking a range of existing self-supervised approaches in molecular representation learning [1][2][3]. It is unconventional to explore scaling laws without considering established pre-training strategies like masking and denoising. 2. The pre-training data in this paper consists of only 5 million molecules. While this scale may be constrained by the label requirements in supervised pre-training, it is insufficient for exploring scaling laws in the molecular domain. Many works use significantly larger datasets, often involving hundreds of millions or even billions of molecules for pre-training their models. 3. The scaling of molecular data appears to perform poorly under the downstream fine-tuning and probing testing protocols. Are there any insights or analyses regarding this phenomenon? 4. From the foundational model perspective, it is also important to evaluate quantum molecular properties such as QM9 and MD17 beyond tasks related to biological ADME properties. [1]:Uni-Mol: a universal 3d molecular representation learning framework [2]:Pre-training Molecular Graph Representation with 3D Geometry [3]:Fractional Denoising for 3D Molecular Pre-training

Questions

See weakness

Rating

3

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

The paper attempts to uncover the significant scaling law challenges in molecular representation learning. However, it is limited by its focus on a single supervised strategy and constrained data scaling. Despite conducting extensive experiments, these efforts fall short of fully addressing the issue and providing actionable conclusions to the community.

Authorsrebuttal2024-08-13

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.

Reviewer FQ727/10 · confidence 4/52024-07-04

Summary

This paper investigates the scaling of GNNs on molecular tasks. In their setting, they pre-train a GNN on multiple molecule dataset and then fine-tune these GNNs on different datasets (“down-stream tasks”). They investigate how the performance of different GNNs changes with parameter count, depth, training samples and so on.

Strengths

- **(S1 - Clarity)** The paper is extremely well written and easy to follow. - **(S2 - Significance)** Both foundation models for graphs and ML on molecules are very active research fields. As such this paper can be important to a large number of people. - **(S3 - Novelty and Significance)** To me, there are four interesting and novel contributions / observations in this paper: 1) Scaling laws for GNNs. 2) The investigation of how to train a pre-trained GNN on a downstream task. Here they investigate two options: fine-tuning the entire model and training only the final MLP (probing). 3) A comparison of three different types of GNNs: MPNNs, Transformers and MPNNs+Transformers. It is particularly interesting that GNNs scale better with parameters than transformers. 4) The observation that removing a dataset from the training data can improve performance on downstream tasks. - **(S4 - Quality)** The evaluation seems solid, the authors compare three different models and observe mostly similar behavior for all models across datasets.

Weaknesses

- **(W1 - Clarity)** The presentation of the results in the main paper is done poorly. Figure 2 is too small and difficult to read / understand. I understand that this paper presents a lot of results but putting all of them into a single figure is in my opinion a bad choice. - **(W2)** Fine tuning vs probing: you have experiments to compare these two approaches. However, I could not find an analysis of the results in the paper (maybe I have missed it).

Questions

- **(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? - **(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? It seems to me like this indicates that when pre-training GNNs the size of the original dataset does not matter much to the downstream task. - **(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)? - **(Q4)** From (W2): how does fine-tuning compare to probing? Is one clearly better than the other? **To sum up,** this is a good experimental paper that should be accepted.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes

Reviewer kaW17/10 · confidence 4/52024-07-13

Summary

The paper examines the scaling behavior of GNNs in the context of molecular graph representation and prediction. The authors analyze various GNN architectures, including message-passing networks, graph Transformers, and hybrid models, on a large dataset of 2D molecular graphs. They find that increasing the scale of model depth, width, dataset size, and diversity significantly enhances performance on downstream tasks. They introduces MolGPS, a new graph foundation model, which demonstrates superior performance across numerous tasks in molecular property prediction.

Strengths

1. The paper is well-written and easy to follow. 2. Comprehensive experiments were performed on large scale datasets to demonstrate the scaling law. 3. The empirical results of MolGPS on various dataset is promising. 4. The insights on how to scale GNN models is meaningful to the community.

Weaknesses

1. Another important aspect of scaling the parameters is regularization [1]. The authors did not explicitly discuss / validate the regularization methods they are using. 2. The authors did not control the compute used to train the models. It would nice to see some efficiency comparison as well. [1] How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers https://arxiv.org/abs/2106.10270

Questions

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

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

NA

Reviewer YrHv6/10 · confidence 3/52024-07-14

Summary

The paper investigates the scalability of GNNs in molecular graph applications. It highlights the relationship between model size, dataset size, and performance upon message-passing networks, graph Transformers, and hybrid architectures using a large dataset of 2D molecular graphs. It shows that supervised pretraining on molecular graphs provides rich embeddings beneficial for downstream tasks. And it finds that the number of labels is crucial for fine-tuning performance.

Strengths

The experiment of MPNN and Transformer on pretrained and downstream tasks is very thorough and provides ample evidence to support the claim. The ablation studies are detailed.

Weaknesses

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. 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.

Questions

See weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors claim that this work uncovers additional aspects of GNN training such as the increasing complexity of aggregation functions and their effect on scaling properties.

Reviewer Lxfs2024-08-12

Thank you for the authors' detailed feedback. However, I still believe that validating the scaling law using the available labeled data may make a trivial contribution to the community. As highlighted in the abstract, "Scaling deep learning models has been at the heart of recent revolutions in language modeling and image generation." The validation of scaling laws in CV and NLP domains largely relies on unlabeled data and self-supervised tasks [1,2,3], as referenced in the paper. I argue that a more relevant approach for studying scaling is under an unlabeled setting, and there are already studies that explore this [4]. In the biological domain, obtaining labeled data is more expensive and challenging compared to the CV or NLP domains. Validating scaling laws in such a constrained setting seems unusual. The effectiveness of this approach may only depend on the relevance of supervised tasks between pre-training and fine-tuning. [1]: OpenAI. Gpt-4 technical report [2]: Llama 2: Open foundation and fine-tuned chat models [3]: Language models are unsupervised multitask learners. [4]: Uni-Mol2: Exploring Molecular Pretraining Model at Scale

Authorsrebuttal2024-08-12

Response to comment by Reviewer Lxfs

We thank the reviewer for their response and are happy to provide further clarification of the concerns mentioned. >**The validation of scaling laws in CV and NLP domains largely relies on unlabeled data and self-supervised tasks [1,2,3] [...]. [A] more relevant approach for studying scaling is under an unlabeled setting [e.g.,] [4].** Unsupervised approaches are an interesting avenue of research in certain areas of the molecular domain. As suggested, Uni-Mol models [4,5] show promising results for modeling molecules in 3D space and for finetuning to physics-based downstream tasks like QM9. We would like to remind the reviewer that we follow a significantly different approach based on modeling 2D graphs with different applications & downstream tasks, where unsupervised models have not yielded comparable results. In the context of CV/NLP [1,2,3] we also point out the lack of equally label-rich datasets that explains why supervised pretraining is not a focal point in those domains. With respect to Uni-Mol2 [4], we note that it has been only finetuned to physics-based tasks (QM9, e.g., homo-lumo gap), that highly depend on 3D coordinates and graph structure. [4] is specialized on understanding graph structure in 3D space, which leads the approach to perform well on such tasks. >We hope that our work paves the way for an era where foundational GNNs drive pharmaceutical drug discovery. [abstract of our submission] Our objectives are different: predicting properties that drive drug discovery like ADMET and binding predictions that strongly rely on understanding the biochemical space beyond graph structure. Uni-Mol2 provides no results on any shared or similar downstream task. While the initial Uni-Mol model [5] provided some results for ADMET downstream tasks, the derivation of the results lacks transparency and is not reproducible (despite our best efforts). Uni-Mol's finetuning compares to GraphMVP [6] (published one year before Uni-Mol), which we also compare to in our work, closely following their experimental setup. Uni-Mol outperforms GraphMVP but reports performance directly taken from [6], despite clearly using a different dataset splitting technique in their work. We reiterate that data splits have a significant impact on empirical results in the context of molecular scaffold splits, especially in the low-sample regime of the MoleculeNet. Further, [5] does not provide sufficient information for reproducibility. We have rigorously analyzed the papers and code of both works and provide details in a separate comment below. Overall, the learnings from our scaling study (width, depth, #molecules, #labels per molecule, composition of pretraining data mixture) lead to the derivation of a foundational GNN that has set a new standard in the various competitive downstream task benchmarks considered here. We kindly request the reviewer to provide more context on their assessment of our work as “a trivial contribution” in light of the above discussion. We would be happy to clarify any further questions the reviewer may have. >**In the biological domain, obtaining labeled data is more expensive and challenging compared to the CV or NLP domains. [...] The effectiveness of this approach may only depend on the relevance of supervised tasks between pre-training and fine-tuning.** We respectfully disagree with this assessment. Adequate labeled data is readily and publicly available at large scale with our pretraining only using a small fraction of the available data. We recall that PCBA_1328 is only a small subset of the PubChem database (only considering datasets with binary tasks with at least 6k molecules) and larger alternatives for PCQM4M exist (e.g., PM6 dataset in [7] that is more than 20x bigger). We also recall that our results show no signs of a data bottleneck. Instead, our molecule scaling suggests, downstream task performance does not deteriorate much when pretraining on a smaller fraction of molecules, e.g., 25% or 50% (Fig. 2 of submission). This relates back to the importance of the number of labels (i.e., label scaling in submission) per molecule already discussed in our initial response. The objective of our work, as in CV/NLP, is to obtain informative embeddings for domain-specific downstream applications. For drug discovery, our work suggests supervised pretraining is a suitable avenue, thanks to the availability of large labeled data sources. As there is no overlap between the pretraining and finetuning tasks, our downstream evaluation solely evaluates the information content of our learned representations (similar to other domains). We are happy to provide more information if further questions arise. [1,2,3,4] as referenced in previous comment of Reviewer Lxfs [5] Uni-Mol [6] Liu et al, Pre-training Molecular Graph Representation with 3D Geometry, ICLR 2022 [7] Beaini et al, Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets, ICLR 2024

Authorsrebuttal2024-08-12

Additional technical details on the inconsistent data splitting in Uni-Mol [5] compared to GraphMVP [6]

Although both papers mention that they use *scaffold* splits, the way they implemented the splitting approach is considerably different, which makes their reported results incomparable. The specific differences are the following: * *GraphMVP*: this model splits the data according to the scaffold splitting approach provided by DeepChem ([here](https://github.com/deepchem/deepchem/blob/64382bd04c9f77f89b10aec3a3351415976ca89a/deepchem/splits/splitters.py#L1481)). The exact code can be found [here](https://github.com/chao1224/GraphMVP/blob/80c42aff21ece462c951cfda110f8eaf866aea0c/src_classification/splitters.py#L19). This function splits the data into train/validation/test splits provided the split ratios. More specifically, the scaffold splitting provided by DeepChem, which is used by GraphMVP and **our work**, is described as follows ([reference](https://deepchem.readthedocs.io/en/latest/api_reference/splitters.html#scaffoldsplitter)): “*ScaffoldSplitter* organizes molecules according to the Bemis-Murcko scaffold representation, which categorizes rings, linkers, frameworks (combinations of linkers and rings), and atomic characteristics like atom type, hybridization, and bond order within a molecular dataset. Next, it sorts these groups by the number of molecules they contain in descending order.” * *Uni-Mol*: this model implements the scaffold splitting according to the [GroupKFold approach provided by sci-kit learn](https://github.com/scikit-learn/scikit-learn/blob/70fdc843a4b8182d97a3508c1a426acc5e87e980/sklearn/model_selection/_split.py#L537). The exact code can be found [here](https://github.com/deepmodeling/Uni-Mol/blob/297587c47a0dd8e8e4c5fc40c15ef6ee821c1440/unimol_tools/unimol_tools/tasks/split.py#L39). More specifically, sci-kit learn GroupKFold splitting approach can be described as follows ([reference](https://scikit-learn.org/stable/modules/cross_validation.html#group-k-fold)): “*GroupKFold* is a variant of k-fold cross-validation designed to prevent the same group from being included in both the training and test sets. GroupKFold helps to identify and avoid such overfitting by ensuring that each group is only present in either the training or test set, but not both.” It is important to note that for the GroupKFold splitting, an important argument is the `groups` argument that specifies the underlying group assignment of the data. In Uni-Mol implementation, this argument is by default retrieved from a file ([here](https://github.com/deepmodeling/Uni-Mol/blob/297587c47a0dd8e8e4c5fc40c15ef6ee821c1440/unimol_tools/unimol_tools/train.py#L32)); the exact file or how it is passed is not provided in the repository. Moreover, we found that (during the fine-tuning) Uni-Mol reads splits from some `.lmbd` files but these files are not provided in the repository (please check [here](https://github.com/deepmodeling/Uni-Mol/blob/297587c47a0dd8e8e4c5fc40c15ef6ee821c1440/unimol/unimol/tasks/unimol_finetune.py#L189) and [here](https://github.com/deepmodeling/Uni-Mol/blob/297587c47a0dd8e8e4c5fc40c15ef6ee821c1440/unimol/unimol/data/lmdb_dataset.py#L18)). It is noteworthy that there are some other important studies, such as GPSE [8] or GraphLoG [9], that also evaluate their results on MoleculeNet with scaffold splitting. However, they either adopt the [OGB splits](https://ogb.stanford.edu/docs/graphprop/), or the sklearn splitting approaches. So, their results are also not directly comparable. [8] Cantürk et al, Graph positional and structural encoder, ICML 2024 [9] Xu et al, Self-supervised graph-level representation learning with local and global structure, ICML 2021

Reviewer kaW12024-08-12

I'm satisfied with the author's response and I'll keep my score.

Reviewer FQ722024-08-12

I thank the authors for answering my questions and providing such a thorough rebuttal. I am still of the opinion that this is a good paper which should be accepted and will thus keep my score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC