[论文解读] 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 和实验跟踪功能，自动化 MLOps 工作流，包括日志记录、检查点保存和超参数跟踪。
• 通过 PyTorch Lightning 的策略模块，以极少的代码更改实现多 GPU 和跨节点的分布式训练，抽象复杂性。

实验结果

研究问题

• RQ1 统一的、模块化的框架是否能降低在 OpenCatalyst 等大规模材料数据集上训练 GNN 的工程开销？
• RQ2 Open MatSci ML Toolkit 如何实现在不同硬件平台和计算基础设施上对称不变 GNN 的一致且可扩展的训练？
• RQ3 使用该框架在 OpenCatalyst 数据集上训练新型 GNN 架构（如 E(n)-GNN 和 MegNet）可达到何种性能？
• RQ4 在分布内（ID）和分布外（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 明显，表明其具有更高的表征能力，并在复杂多组分任务上优化效果更优。