[논문 리뷰] The Open MatSci ML Toolkit: A Flexible Framework for Machine Learning in Materials Science

Santiago Miret, Kin Long Kelvin Lee외 3명|arXiv (Cornell University)|2022. 10. 31.

Machine Learning in Materials Science인용 수 6

한 줄 요약

Open MatSci ML 툴킷은 PyTorch Lightning과 DGL 기반으로 구축된 유연하고 오픈소스인 파이썬 프레임워크로, 재료 과학 분야에서 확장 가능하고 재현 가능한 기계학습 워크플로우를 가능하게 합니다. 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 및 실험 추적 기능을 통해 로깅, 체크포인팅, 하이퍼파rameter 추적 등 MLOps 워크플로우를 자동화합니다.
• 최소한의 코드 변경으로 다중 GPU 및 노드 간 분산 학습을 가능하게 하며, PyTorch Lightning의 전략 모듈을 통해 복잡성을 추상화합니다.

실험 결과

연구 질문

• RQ1 통합적이고 모듈러한 프레임워크는 OpenCatalyst와 같은 대규모 재료 과학 데이터셋에서 GNN을 학습할 때 발생하는 기술적 부담을 줄일 수 있는가?
• RQ2 Open MatSci ML 툴킷은 다양한 하드웨어 플랫폼과 컴퓨팅 인fra구조에서 등방성 GNN의 일관되고 확장 가능한 학습을 어떻게 지원하는가?
• RQ3 이 프레임워크를 사용해 OpenCatalyst 데이터셋에서 E(n)-GNN 및 MegNet과 같은 새로운 GNN 아키텍처를 학습했을 때 기대할 수 있는 성능은 어떠한가?
• RQ4 특히 새로운 흡착제에 대해 예측할 때, 분포 내(ID)와 분포 외(OOD) 설정 간의 모델 일반화 능력은 어떻게 비교되는가?
• RQ5 이 툴킷은 원자 중심 포인트 클러스터를 사용하는 등, 파ram터 효율적이고 도메인 특화된 모델 설계를 어느 정도 지원할 수 있는가?

주요 결과

• Open MatSci ML 툴킷은 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 작업는 IS2RE에 비해 학습 도중 손실 감소 폭이 더 뚜렷했으며, 특히 MegNet의 경우 더 높은 표현 능력과 복잡한 다성분 작업에 대한 더 나은 최적화 능력을 보였습니다.