Geometry-enhanced Pre-training on Interatomic Potentials

Abstract: Machine learning interatomic potentials (MLIPs) enables molecular dynamics (MD) simulations with ab initio accuracy and has been applied to various fields of physical science. However, the performance and transferability of MLIPs are limited by insufficient labeled training data, which require expensive ab initio calculations to obtain the labels, especially for complex molecular systems. To address this challenge, we design a novel geometric structure learning paradigm that consists of two stages. We first generate a large quantity of 3D configurations of target molecular system with classical molecular dynamics simulations. Then, we propose geometry-enhanced self-supervised learning consisting of masking, denoising, and contrastive learning to better capture the topology and 3D geometric information from the unlabeled 3D configurations. We evaluate our method on various benchmarks ranging from small molecule datasets to complex periodic molecular systems with more types of elements. The experimental results show that the proposed pre-training method can greatly enhance the accuracy of MLIPs with few extra computational costs and works well with different invariant or equivariant graph neural network architectures. Our method improves the generalization capability of MLIPs and helps to realize accurate MD simulations for complex molecular systems.

• The paper introduces a novel framework combining classical molecular dynamics and self-supervised learning to enhance MLIPs for accurate energy and force predictions.
• It incorporates three geometry-enhanced SSL tasks that capture topological and spatial information to improve the understanding of molecular structures.
• Experimental results demonstrate significant accuracy gains and reduced computation costs, paving the way for more scalable and robust molecular simulations.

Geometry-enhanced Pre-training on Interatomic Potentials

The paper "Geometry-enhanced Pre-training on Interatomic Potentials" presents an innovative framework for advancing machine learning interatomic potentials (MLIPs) using a geometry-enhanced approach to address limitations in data scarcity and computation cost. MLIPs are pivotal in simulating molecular dynamics (MD) by accurately modeling interatomic interactions. However, their efficacy largely depends on high-quality labeled data derived from expensive ab initio calculations, which limits their scalability and application range.

Framework Overview

The proposed framework consists of two primary stages: the generation of unlabeled molecular configurations using classical molecular dynamics (CMD) simulations with empirical force fields, and a subsequent pre-training phase utilizing self-supervised learning (SSL) techniques. The CMD simulations exploit empirical interatomic potentials, which allow for efficient generation of atomic configurations albeit without the precision required for directly calculating energies and forces.

In the pre-training phase, the framework introduces three geometry-enhanced SSL tasks:

1. Masked autoencoder reconstruction of atom features with added noisy coordinates to learn topological information.
2. Denoising tasks that predict noise added to atomic positions, enforcing the model to recover spatial information.
3. Contrastive learning utilizing a 3D network to align representations learned through global geometry, improving the model's capture of three-dimensional structural nuances.

These tasks collectively aim to refine the geometric understanding of molecular systems, thereby enhancing the generalizability and performance of MLIPs with minimal additional computational cost compared to utilizing larger labeled datasets.

Experimental Evaluation

The framework was tested using various data sets, including MD17 and ISO17, which involve small organic molecules, as well as datasets for liquid water and electrolyte solutions with periodic boundary conditions (PBCs). The results showed that the introduction of geometry-enhanced pre-training significantly improved the accuracy of force and energy predictions across multiple MLIP architectures, including SchNet, DimeNet, and SphereNet. Notably, the performance gains were achieved with a remarkable reduction in the computation time required for generating labeled data through ab initio methods.

Benchmark Performance

GPIP-based models demonstrated high accuracy, dynamic stability, and robustness in predicting molecular properties. For instance, the SchNet-G model, enhanced by SSL tasks, exhibited a significant reduction in mean absolute error (MAE) in both energy and force predictions compared to its baseline. This improvement was consistent across extensive benchmarks, indicative of the framework's adaptability to various types of molecular systems.

Implications and Future Work

The ability to leverage CMD-generated datasets, scalable across diverse chemical domains, introduces a new paradigm in dataset generation that can significantly lower the entry barrier for high-fidelity molecular simulations. The minor computational overhead of GPIP, relative to traditional data generation approaches, proposes a sustainable alternative for training data acquisition.

The study opens avenues for future exploration in self-supervised molecular representations, particularly towards integrating more complex geometric transformations or incorporating additional physical constraints into the learning paradigm. Furthermore, extending this approach to other molecular simulation challenges, such as multiscale modeling or reactive dynamics, could broaden the applicability of MLIPs further.

In summary, by combining CMD with geometry-enhanced SSL, the paper demonstrates a substantial advance in the efficiency and efficiency of MLIPs, providing a robust groundwork for further exploration and refinement of ML-based molecular simulations.