LaMM: Semi-Supervised Pre-Training of Large-Scale Materials Models

TL;DR

LaMM introduces a semi-supervised pre-training approach for neural network potentials in materials science, leveraging large-scale semi-labeled data and load balancing to enhance training efficiency and model accuracy.

Contribution

The paper presents LaMM, a novel semi-supervised pre-training method with improved denoising and load balancing, enabling efficient training on large semi-labeled datasets for materials modeling.

Findings

Effective utilization of 300 million semi-labeled samples

Improved fine-tuning speed and accuracy

Enhanced load balancing during multi-node training

Abstract

Neural network potentials (NNPs) are crucial for accelerating computational materials science by surrogating density functional theory (DFT) calculations. Improving their accuracy is possible through pre-training and fine-tuning, where an NNP model is first pre-trained on a large-scale dataset and then fine-tuned on a smaller target dataset. However, this approach is computationally expensive, mainly due to the cost of DFT-based dataset labeling and load imbalances during large-scale pre-training. To address this, we propose LaMM, a semi-supervised pre-training method incorporating improved denoising self-supervised learning and a load-balancing algorithm for efficient multi-node training. We demonstrate that our approach effectively leverages a large-scale dataset of $\sim$300 million semi-labeled samples to train a single NNP model, resulting in improved fine-tuning performance in terms of both speed and accuracy.