[论文解读] On the Scalability of GNNs for Molecular Graphs

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

Machine Learning in Materials Science被引用 6

一句话总结

本文研究不同 GNN 架构在大分子图上的可伸缩性，显示更大的宽度、深度、数据量以及多样化的预训练带来持续的收益，并提出 MolGPS 作为在多项任务上达到SOTA的图型基础模型。

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.

研究动机与目标

• 评估宽度、深度、数据集大小、标签数量和数据集多样性如何影响分子图上的 GNN 性能。
• 在预训练、微调和探针设置中对比三种架构（MPNN++, Graph Transformer, GPS++）。
• 确定构建分子图基础模型的有效预训练策略和架构选择。
• 为分子发现提供迈向 MolGPS 风格的基础模型的实际指导。

提出的方法

• 使用 Graphium 库研究三种架构的尺度行为：MPNN++, Graph Transformer 和 GPS++。
• 在 LargeMix 上进行预训练，这是一个包含多任务的 大型数据集，最多有 5 million 个分子和数千个标签。
• 在 TDC、Polaris 和 MoleculeNet 的 38 个下游任务上评估微调和探针。
• 使用位置编码（PSEs）和结构偏置来增强图变换器。
• 分析相对于宽度、深度、数据集大小和标签的尺度；将观察到的趋势拟合为幂律尺度关系（L_theta ~ (D_C / D)^beta 和 L_theta ~ (theta_C / theta)^alpha）。
• 通过在不使用 L1000 的前提下，将 MPNN++、Transformer 和 GPS++ 的指纹组合起来提出 MolGPS，随后添加一个 MLP 头。

实验结果

研究问题

• RQ1 当增加宽度、深度、数据集大小和标签多样性时，分子 GNN 的尺度如何变化？
• RQ2 哪种架构（MPNN++, Graph Transformer, GPS++）在分子性质预测和下游任务中扩展性最好？
• RQ3 哪些预训练和微调/探针策略能在 38 个下游分子基准上最大化性能？
• RQ4 一个基础模型（MolGPS）能否统一架构选择和预训练以在各任务上获得最先进的结果？

主要发现

• 所有模型在增加宽度、深度和数据集规模时均呈现有利的扩展。
• MPNN++ 在较小/数据稀缺任务上参数高效且数据高效；Transformer 更需要数据但从宽度获益并能在更大数据下扩展。
• 混合型 GPS++ 在低参数条件下受益于 MPNN++，在较大数据条件下的扩展性与 Transformer 类似。
• 增加数据集规模（分子）可带来显著的性能提升，在量子任务上 GPS++ 和 Transformer 获得的提升更多。
• 探针结果表明 Transformer 在下游任务中受益于宽度；深度尺度呈现混合趋势；标签尺度对性能有显著提升。
• 提出 MolGPS（结合 MPNN++、Transformer、GPS++ 的指纹）在 38 个下游任务中有 26 项达到 SOTA，并推动了分子发现领域的统一图基础模型。
• 在没有 L1000 任务的预训练即可在各尺度上提升下游性能；多指纹探针进一步提升结果。