NeurIPS 2024

On the Scalability of GNNs for Molecular Graphs

Maciej Sypetkowski, Frederik Wenkel, Farimah Poursafaei, Nia Dickson, Karush Suri, Philip Fradkin, Dominique Beaini

Main figure for On the Scalability of GNNs for Molecular Graphs

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.

Benchmark Results

Tested on 4 datasets 1 baseline value reported too low

Tables show presented results in global context (best 3 papers + GCN + MLP).

MoleculeNet

#ModelROC-AUC ↑Paper
1
ProtoMol
0.9140 ±0.0030
2
MMSG
0.9080 ±0.0050
3
MemGNN
0.9070
43
GIN
0.8646 ±0.0081
104
3B MolGPS
0.8320
This paper
111
1B MolGPS
0.8280
This paper
#ModelROC-AUC ↑Paper
1
MAGPrompt+
0.9432 ±0.0238
2
CamS-LLaMA
0.9420
3
SE(3)-I MPPH
0.9400 ±0.0210
34
GIN
0.8875 ±0.0049
59
3B MolGPS
0.8090
This paper
60
1B MolGPS
0.8000
This paper
#ModelROC-AUC ↑Paper
1
SE(3)-I MPPH
0.9930 ±0.0040
2
KA-GCN
0.9920 ±0.0050
3
KA-GAT
0.9910 ±0.0050
19
GCN
0.9130 ±0.0170
49
GraphMVP
actual #49 vs paper #108 59 places worse
0.8650 ±0.0170
best reported
actual #49 vs paper #108 59 places worse
paper 0.7750
81
1B MolGPS
0.8090
This paper
83
3B MolGPS
0.8070
This paper
SOTA performance (not significantly different from the best) Proposed method (this paper) Global context (other papers) Best at pub — reported result worse than the best known when published (best reported in green)