[論文レビュー] 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性能にどう影響するかを評価する。
• 事前学習、ファインチューニング、およびプロービング設定にわたり、3つのアーキテクチャ（MPNN++, Graph Transformer, GPS++）を比較する。
• 分子のグラフ基盤モデルを構築するための効果的な事前学習戦略とアーキテクチャ選択を決定する。
• 分子発見のための MolGPSスタイルの基盤モデルへ向けた実践的指針を提供する。

提案手法

• Graphiumライブラリを用いて、3つのアーキテクチャ（MPNN++, Graph Transformer, GPS++）のスケーリング挙動を研究する。
• LargeMixという、最大5百万の分子と数千のラベルを含む大規模マルチタスクデータセットで事前学習を行う。
• 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++はパラメータ効率とデータ効率が高く、小規模/データ不足タスクで優れる。トランスフォーマーはデータ集約性を要するが、幅の恩恵を受け、より大規模データでスケールできる。
• ハイブリッドGPS++は低パラメータ領域でMPNN++の利点を活かし、より大規模データ領域ではトランスフォーマーと同様にスケールする。
• データセットサイズ（分子数）を増やすと顕著な性能向上を得られ、GPS++とトランスフォーマーは量子タスクでより多くの利得を得る。
• プロービングの結果、トランスフォーマーは下流タスクで幅の恩恵を受ける。一方、深さのスケーリングは傾向がまちまちで、ラベルスケーリングは性能を強く向上させる。
• MPNN++, Transformer, GPS++のフィンガープリントを結合したMolGPSの提案は、38の下流タスクのうち26件でSOTAを達成し、分子発見のための統一されたグラフ基盤モデルを推進する。
• L1000タスクなしの事前学習はスケールを問わず下流性能を向上させ、複数フィンガープリントのプロービングはさらに結果を向上させる。