[Paper Review] On the Scalability of GNNs for Molecular Graphs

Maciej Sypetkowski, Frederik Wenkel+5 more|arXiv (Cornell University)|Apr 17, 2024

Machine Learning in Materials Science6 citations

TL;DR

The paper studies how different GNN architectures scale for large molecular graphs, showing that larger width, depth, data, and diversified pretraining yield continued gains, and introduces MolGPS as a graph foundation model achieving SOTA on many tasks.

ABSTRACT

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.

Motivation & Objective

• Assess how width, depth, dataset size, label count, and dataset diversity affect GNN performance on molecular graphs.
• Compare three architectures (MPNN++, Graph Transformer, GPS++) across pretraining, finetuning, and probing settings.
• Determine effective pretraining strategies and architecture choices for building a graph foundation model for molecules.
• Provide practical guidance toward a MolGPS-style foundation model for molecular discovery.

Proposed method

• Study scaling behavior of three architectures: MPNN++, Graph Transformer, and GPS++ using the Graphium library.
• Pretrain on LargeMix, a large multi-task dataset with up to 5 million molecules and thousands of labels.
• Evaluate finetuning and probing on 38 downstream tasks from TDC, Polaris, and MoleculeNet.
• Use positional encodings (PSEs) and structural biases to enhance graph transformers.
• Analyze scaling with respect to width, depth, dataset size, and labels; fit observed trends to power-law scaling relations (L_theta ~ (D_C / D)^beta and L_theta ~ (theta_C / theta)^alpha).
• Propose MolGPS by combining fingerprints from MPNN++, Transformer, and GPS++ after pretraining without L1000, followed by an MLP head.

Experimental results

Research questions

• RQ1 How do molecular GNNs scale when increasing width, depth, dataset size, and label diversity?
• RQ2 Which architectures (MPNN++, Graph Transformer, GPS++) scale best for molecular property prediction and downstream tasks?
• RQ3 What pretraining and fine-tuning/probing strategies maximize performance on 38 downstream molecular benchmarks?
• RQ4 Can a foundation model (MolGPS) unify architectural choices and pretraining to achieve state-of-the-art results across tasks?

Key findings

• All models scale favorably with increased width, depth, and dataset size.
• MPNN++ is more parameter-efficient and data-efficient on smaller/Data-scarce tasks; Transformers are more data-hungry but benefit from width and can scale with larger data.
• Hybrid GPS++ benefits from MPNN++ in low-parameter regimes and scales similarly to Transformers in larger data regimes.
• Increasing dataset size (molecules) yields substantial performance gains, with GPS++ and Transformer gaining more on quantum tasks.
• Probing shows Transformers benefit from width for downstream tasks; depth scaling shows mixed trends; label scaling strongly boosts performance.
• Proposing MolGPS (combining fingerprints from MPNN++, Transformer, GPS++) yields SOTA on 26 of 38 downstream tasks and advances a unified graph foundation model for molecular discovery.
• Pretraining without L1000 tasks improves downstream performance across scales; multi-fingerprint probing further enhances results.