DeePMD-kit v3: A Multiple-Backend Framework for Machine Learning Potentials

TL;DR

DeePMD-kit v3 introduces a multi-backend framework supporting TensorFlow, PyTorch, JAX, and PaddlePaddle, enabling seamless integration and interoperability of machine learning potentials across different scientific applications.

Contribution

This work presents a multi-backend architecture for DeePMD-kit, allowing flexible backend switching and integration with various machine learning frameworks, enhancing usability and extensibility.

Findings

Supports multiple ML frameworks including TensorFlow, PyTorch, JAX, PaddlePaddle

Enables seamless backend switching with minimal modifications

Facilitates integration of different MLP packages and workflows

Abstract

In recent years, machine learning potentials (MLPs) have become indispensable tools in physics, chemistry, and materials science, driving the development of software packages for molecular dynamics (MD) simulations and related applications. These packages, typically built on specific machine learning frameworks such as TensorFlow, PyTorch, or JAX, face integration challenges when advanced applications demand communication across different frameworks. The previous TensorFlow-based implementation of DeePMD-kit exemplified these limitations. In this work, we introduce DeePMD-kit version 3, a significant update featuring a multi-backend framework that supports TensorFlow, PyTorch, JAX, and PaddlePaddle backends, and demonstrate the versatility of this architecture through the integration of other MLPs packages and of Differentiable Molecular Force Field. This architecture allows seamless backend switching with minimal modifications, enabling users and developers to integrate DeePMD-kit with other packages using different machine learning frameworks. This innovation facilitates the development of more complex and interoperable workflows, paving the way for broader applications of MLPs in scientific research.