[論文レビュー] The Open MatSci ML Toolkit: A Flexible Framework for Machine Learning in Materials Science

Santiago Miret, Kin Long Kelvin Lee+3 more|arXiv (Cornell University)|Oct 31, 2022

Machine Learning in Materials Science被引用数 6

ひとこと要約

Open MatSci ML Toolkit は、PyTorch Lightning と DGL に基づいた柔軟でオープンソースの Python フレームワークであり、材料科学におけるスケーラブルで再現可能な機械学習ワークフローを可能にします。OpenCatalyst データセット上で等長性グラフニューラルネットワークの学習を高速化し、優れた性能を達成しています—例えば、OOD-Both S2EF タスクで 0.38 MAE を達成—同時に CPU、GPU、クラスタ間でのデプロイを簡素化しています。

ABSTRACT

We present the Open MatSci ML Toolkit: a flexible, self-contained, and scalable Python-based framework to apply deep learning models and methods on scientific data with a specific focus on materials science and the OpenCatalyst Dataset. Our toolkit provides: 1. A scalable machine learning workflow for materials science leveraging PyTorch Lightning, which enables seamless scaling across different computation capabilities (laptop, server, cluster) and hardware platforms (CPU, GPU, XPU). 2. Deep Graph Library (DGL) support for rapid graph neural network prototyping and development. By publishing and sharing this toolkit with the research community via open-source release, we hope to: 1. Lower the entry barrier for new machine learning researchers and practitioners that want to get started with the OpenCatalyst dataset, which presently comprises the largest computational materials science dataset. 2. Enable the scientific community to apply advanced machine learning tools to high-impact scientific challenges, such as modeling of materials behavior for clean energy applications. We demonstrate the capabilities of our framework by enabling three new equivariant neural network models for multiple OpenCatalyst tasks and arrive at promising results for compute scaling and model performance.

研究の動機と目的

• OpenCatalyst のような大規模な材料科学データセットに機械学習を適用する際の高い工学的・使いやすさの障壁を解消すること。
• OpenCatalyst データセット上でモデルの学習とデプロイを可能にする、自己完結的でモジュラーなフレームワークを提供することで、新規研究者の参入障壁を低減すること。
• 科学的抽象化を損なわず、ラップトップからクラスタまで多様なハードウェア（CPU、GPU、XPU）およびコンピューティング環境での実験のスケーリングを円滑に可能にすること。
• DGL を用いたグラフニューラルネットワークの迅速なプロトタイピングを可能にし、OpenCatalyst エコシステムにおける既存の PyG ベースのツールチェーンを補完すること。
• E(n)-GNN や MegNet などの新しい幾何的ディープラーニングアーキテクチャの開発とベンチマークを促進し、材料物性予測に応用すること。

提案手法

• PyTorch Lightning を活用し、CPU、GPU、クラスタ間での分散学習、混合精度学習、ハードウェアスケーリングを抽象化することで、学習の管理を簡素化。
• DGL を統合し、効率的なグラフニューラルネットワーク（GNN）モデル開発を可能にし、分子および結晶グラフ上でメッセージパッシングアーキテクチャの迅速なプロトタイピングを実現。
• OpenCatalyst データセット用の事前構築済みデータローダーを備えたモジュラーで自己完結的なコードベースを提供し、ローカル実験用に開発用データセットを含む。
• 原子中心の点群表現を用いた、IS2RE（吸着エネルギー予測）および S2EF（エネルギーおよび力予測）の複数タスクに対するエンドツーエンドの学習パイプラインをサポート。
• PyTorch Lightning の組み込み CLI および実験追跡機能を活用し、ログ記録、チェックポイント保存、ハイパーパramータ追跡の MLOps ワークフローを自動化。
• 最小限のコード変更でマルチGPUおよびノード間分散学習を可能にし、PyTorch Lightning の戦略モジュールによって複雑さを抽象化。

実験結果

リサーチクエスチョン

• RQ1 統合的でモジュラーなフレームワークは、OpenCatalyst のような大規模な材料科学データセットにおける GNN 学習の工学的負担を低減できるか？
• RQ2 Open MatSci ML Toolkit は、多様なハードウェアプラットフォームおよびコンピューティングインfraストラクチャ間で、等長性 GNN の一貫性とスケーラビリティをどのように実現するか？
• RQ3 このフレームワークを用いて、E(n)-GNN や MegNet などの新しい GNN アーキテクチャを OpenCatalyst データセットで学習した場合、どの程度の性能が達成できるか？
• RQ4 未知の吸着剤を含む分布外（OOD）設定において、モデルの一般化能力はどのように比較されるか、特に未確認の吸着剤の予測が困難である点に注目。
• RQ5 このツールキットは、原子中心の点群を用いたパラメータ効率的でドメイン特化型のモデル設計をどの程度サポートできるか？

主な発見

• Open MatSci ML Toolkit は、OpenCatalyst データセット上で等長性 GNN の学習を最小限のセットアップで可能にし、ラップトップでのローカル開発からクラスタでの大規模分散学習までをカバー。
• MegNet は S2EF タスクにおいて OOD-Both のテスト誤差を 0.38 に達成し、はるかに大きなデータセット（133M サンプル）で学習された GemNet-XL と同等の性能を示した。
• E(n)-GNN は OOD-Both S2EF タスクでテスト誤差 0.43 を達成し、MegNet よりもパラメータ数が少ないにもかかわらず、強力な性能を示した。
• すべてのモデルが OOD-Adsorbates への一般化に困難を示し、新規分子の物性を予測することが依然として主要な課題であることが示された。
• Gala モデルは ID および OOD タスクの両方で一貫した性能を示し、E(n)-GNN や MegNet よりも優れた一般化能力を有していることが示唆された。
• S2EF タスクでは、特に MegNet において、IS2RE よりも学習過程での損失低下が急激に進んだ。これは、複雑で多成分のタスクにおいて、より高い表現能力と優れた最適化能力を有していることを示している。