[論文レビュー] Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets

Dominique Beaini, Shenyang Huang+33 more|arXiv (Cornell University)|Jan 1, 2023

Machine Learning in Materials Science被引用数 6

ひとこと要約

本稿では、約1億個の分子と3,000のタスクにわたる130億以上のラベル付き量子的・生物学的性質を有する、大規模でマルチタスクな分子データセットの新ファミリーであるToyMix、LargeMix、UltraLargeを紹介する。著者らは、効率的なマルチタスクおよびマルチレベルのグラフ学習を目的とした深層学習ライブラリ「Graphium」を提案し、多様な教師ありデータを用いた事前学習が、リソースが限られた生物学的タスクにおける性能向上を実現することを示した。これは、分子AI分野における基礎的モデルの実現可能性を支持する。

ABSTRACT

Recently, pre-trained foundation models have enabled significant advancements in multiple fields. In molecular machine learning, however, where datasets are often hand-curated, and hence typically small, the lack of datasets with labeled features, and codebases to manage those datasets, has hindered the development of foundation models. In this work, we present seven novel datasets categorized by size into three distinct categories: ToyMix, LargeMix and UltraLarge. These datasets push the boundaries in both the scale and the diversity of supervised labels for molecular learning. They cover nearly 100 million molecules and over 3000 sparsely defined tasks, totaling more than 13 billion individual labels of both quantum and biological nature. In comparison, our datasets contain 300 times more data points than the widely used OGB-LSC PCQM4Mv2 dataset, and 13 times more than the quantum-only QM1B dataset. In addition, to support the development of foundational models based on our proposed datasets, we present the Graphium graph machine learning library which simplifies the process of building and training molecular machine learning models for multi-task and multi-level molecular datasets. Finally, we present a range of baseline results as a starting point of multi-task and multi-level training on these datasets. Empirically, we observe that performance on low-resource biological datasets show improvement by also training on large amounts of quantum data. This indicates that there may be potential in multi-task and multi-level training of a foundation model and fine-tuning it to resource-constrained downstream tasks.

研究の動機と目的

• 分子機械学習における基礎的モデルを訓練するための、大規模でラベル付きかつマルチタスクな分子データセットの不足を解消すること。
• 量子力学的および実験的生物学的アッセイから得られるマルチモodalな教師ありラベルを導入することで、自己教師あり事前学習の限界を克服すること。
• Graphiumという深層学習ライブラリの開発を通じて、大規模で多様な分子データセットに対する効率的な学習を可能にすること。
• マルチタスクおよびマルチレベルの事前学習が、リソースが限られた下流タスクにおける性能向上に寄与することを示す強力なベースラインモデルと実証的証拠を確立すること。

提案手法

• 密度汎関数理論（DFT）（例：B3LYP）および半経験的（例：PM6）手法を用いて計算された量子力学的（QM）性質を統合することで、既存の分子データセットを精査・強化した。
• 高スループットアッセイから得られた生物学的活性ラベル（例：投与量応答プロファイル、遺伝子発現、毒性データ）を収集し、ノードレベルおよびグラフレベルの両方のラベルを含むマルチレベルのラベルを構築した。
• 小規模からPubChemのほぼ完全カバレッジに至るまでをカバーする3つのデータセットカテゴリ（ToyMix、LargeMix、UltraLarge）を設計し、合計で130.4億のラベルを実現した。
• PyTorchベースのライブラリ「Graphium」を開発し、マルチタスクおよびマルチレベルのグラフ学習に最適化されたもので、混合精度学習、モデルパイプライン処理、分散推論をサポートしている。
• メッセージパッシングニューラルネットワークおよびトランスフォーマーを用いたベースラインモデルを実装し、全データセット階層上で学習することで、転移学習性能を評価した。
• 量子的および生物学的タスクの両方で回帰および分類のターゲットを組み合わせ、情報量を最大化し、効果的な事前学習を可能にした。

実験結果

リサーチクエスチョン

• RQ1 量子的および生物学的両方のラベルを有する大規模でマルチタスク・マルチレベルの分子データセットは、分子機械学習における基礎的モデルの効果的な事前学習を可能にするか？
• RQ2 自己教師ありまたは単一タスクの事前学習と比較して、多様な教師ありデータを用いた事前学習が、リソースが限られた生物学的性質予測タスクの性能向上に寄与するか？
• RQ3 マルチタスクおよびマルチレベルの学習目的が、分子モデリングタスク全体における汎化性および転送性をどの程度向上させるか？
• RQ4 OGB-LSC や QM1B といった既存ベンチマークと比較して、提案されたデータセットのデータ量およびラベルの豊かさ（ラベルの豊富さ）はどの程度か？
• RQ5 GPT-2 などの基礎的NLPモデルが使用する事前学習データの規模に近い、統合された深層学習ライブラリ（例：Graphium）が、複数のハードウェアプラットフォーム上でこのような大規模で多様な分子データセットの学習および推論を効果的にサポートできるか？

主な発見

• 提案されたデータセットは、広く使われているOGB-LSCのPCQM4Mv2データセットの300倍以上、量子のみのデータセットであるQM1Bの13倍以上のデータポイントを有する。
• データセットは約1億個の分子と3,000の疎に定義されたタスクをカバーし、量子的および生物学的性質の合計で130億以上の個別のラベルを含む。
• ベースライン結果から、量子データの大量を含む事前学習を経た後、リソースが限られた生物学的データセットでの微調整が顕著に向上することが示された。これは、強力な転移学習の可能性を示している。
• Graphiumライブラリは、大規模なマルチタスクデータセットにおける効率的な学習を可能にし、複数のアクセラレータを活用した混合精度学習および分散推論をサポートしている。
• 提案されたデータセットに含まれるラベル数（130.4億）は、GPT-2のような基礎的NLPモデルが使用する事前学習データのスケールに近づいており、分子表現学習における同等の事前学習容量を有すると示唆している。
• 量子的および生物学的性質の統合的モデリングは、モデルの汎化性を向上させ、多様で教師ありの事前学習が、効果的な分子基礎的モデルを構築するために不可欠であるという仮説を支持している。