[論文レビュー] From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction

Nima Shoghi, Adeesh Kolluru+4 more|arXiv (Cornell University)|Oct 25, 2023

Machine Learning in Materials Science被引用数 21

ひとこと要約

この研究は複数の異なる化学領域に跨る単一モデルを訓練する Joint Multi-domain Pre-training (JMP) を導入し、ゼロから訓練する場合と比較して平均59%の改善を達成し、40課題中34課題で最先端と同等またはそれを上回る。

ABSTRACT

Foundation models have been transformational in machine learning fields such as natural language processing and computer vision. Similar success in atomic property prediction has been limited due to the challenges of training effective models across multiple chemical domains. To address this, we introduce Joint Multi-domain Pre-training (JMP), a supervised pre-training strategy that simultaneously trains on multiple datasets from different chemical domains, treating each dataset as a unique pre-training task within a multi-task framework. Our combined training dataset consists of $\sim$120M systems from OC20, OC22, ANI-1x, and Transition-1x. We evaluate performance and generalization by fine-tuning over a diverse set of downstream tasks and datasets including: QM9, rMD17, MatBench, QMOF, SPICE, and MD22. JMP demonstrates an average improvement of 59% over training from scratch, and matches or sets state-of-the-art on 34 out of 40 tasks. Our work highlights the potential of pre-training strategies that utilize diverse data to advance property prediction across chemical domains, especially for low-data tasks. Please visit https://nima.sh/jmp for further information.

研究の動機と目的

• 多様な化学領域に跨る原子性質予測の一般化の課題を動機づけ、これに対処する。
• 大規模で異種混在のデータセットを活用するスケーラブルな事前訓練戦略を開発する。
• 限られたデータでの下流タスクのファインチューニングを可能にしつつ高い性能を維持する。
• 事前訓練ドメインを超えた未知の領域（大分子および材料）への転移性を示す。

提案手法

• Joint Multi-domain Pre-training (JMP) をマルチタスク監督付き事前訓練フレームワークとして提案する。
• 単一のバックボーンモデル (GemNet-OC) を用い、データセットごとにエネルギーと力の予測ヘッドを設ける。
• データセットごとに線形エネルギー参照と力の正規化を単位ガウス分布へ正規化してターゲットを正規化する。
• バッチ構築時のデータセットサイズのバランスを取るため温度ベースのサンプリングを適用する。
• 異なる系サイズを持つデータセットからの寄与をバランスさせるため構造別の損失削減を導入する。
• 正則化（ウェイト減衰、エッジドロップアウト、EMA）を伴うマルチタスク損失の単位スカラー化を採用する。
• 事前訓練ヘッドをタスク特化ヘッドに置換してファインチューニングを行い、必要に応じてエネルギー勾配から力を計算する。

実験結果

リサーチクエスチョン

• RQ1 複数の化学ドメインで訓練された単一の事前訓練モデルは、小分子・大分子・材料といった下流タスク全般へどれだけ一般化できるか。
• RQ2 ジョイント多ドメイン事前訓練は、スクラッチからの訓練や従来の単一ドメインや自己教師付きアプローチより、多様なベンチマークで上回るか。
• RQ3 データのバランス、損失設計、正則化戦略がマルチタスク事前訓練の性能に与える影響は何か。
• RQ4 JMP は限定的な下流データで大規模モデルのファインチューニングを可能にし、未知の領域への転移を改善できるか。

主な発見

• JMP はファインチューニングタスクでスクラッチ訓練と比較して平均59%の改善をもたらす。
• JMP は QM9、rMD17、MD22、SPICE、MatBench、QMOF の40件中34件のファインチューニングタスクで現状最先端と同等またはそれを上回る。
• 235M パラメータの JMP モデルが複数の低データベースで最先端の性能を達成。
• ファインチューニング JMP-L は GN-OC-L の性能に約1/12の訓練時間で到達し、適応が速いことを示す。
• 多様な化学データでの事前訓練は、非平衡配置やドメイン外ターゲット（例：MatBenchとQMOFの材料特性）へ一般化する転移可能な表現を提供する。
• 全体の JMP 事前訓練コストは、スクラッチ訓練と比較して下流のファインチューニングが12倍以上速くなることによって相殺される。