Introduction

Paper link: https://arxiv.org/abs/2312.15492

The codes, datasets and input scripts are all available on zenodo (https://doi.org/10.5281/zenodo.10428497)

This notebook is updated on 9.14 based on beta3 code version.

The rapid development of artificial intelligence (AI) is driving significant changes in the field of atomic modeling, simulation, and design. AI-based potential energy models have been successfully used to perform large-scale and long-time simulations with the accuracy of ab initio electronic structure methods. However, the model generation process still hinders applications at scale. We envision that the next stage would be a model-centric ecosystem, in which a large atomic model (LAM), pre-trained with as many atomic datasets as possible and can be efficiently fine-tuned and distilled to downstream tasks, would serve the new infrastructure of the field of molecular modeling. We propose DPA-2, a novel architecture for a LAM, and develop a comprehensive pipeline for model fine-tuning, distillation, and application, associated with automatic workflows. We show that DPA-2 can accurately represent a diverse range of chemical systems and materials, enabling high-quality simulations and predictions with significantly reduced efforts compared to traditional methods. Our approach paves the way for a universal large atomic model that can be widely applied in molecular and material simulation research, opening new opportunities for scientific discoveries and industrial applications.

In order to run this notebook successfully, let's do some preparatory work first.

• data: This directory contains three datasets FerroEle, H2O-PD and SemiCond. FerroEle is a small subset of the dataset FerroEle_DPA_v1_0. H2O_H2O-PD is a small subset of the dataset H2O-PD_DPA_v1_0. SemiCond is a small subset of the dataset SemiCond_DPA_v1_0.If you want to get the full dataset, you can go to website AIS Square to download these datasets.
• model: This directory contains one singletask model H2O-PD.pt and one multitask model OpenLAM_2.2.0_27heads_beta3.pt. The multitask model OpenLAM_2.2.0_27heads_beta3.pt is trained on 27 different datasets.
• train: This directory contains three modes of training, which are singletask training, multitask training, and finetuning based on pretrained model. In a moment, we're going to demonstrate how to do these trainings in the three folders.
• md: In this directory, we will demonstrate how to perform molecular dynamics simulation using DPA-2 model.

Model Loading

One can download someone else's trained singletask model or multitask model from web AIS Square

command line interface

The model can be used in many ways. The most straightforward test can be performed using dp test

Firstly we test the valid dataset of H2O_H2O-PD using the singletask model H2O-PD.pt.

where -m gives the model checkpoint to import, -s the path to the tested system and -n the number of tested frames. Several other command line options can be passed to dp test, which can be checked with

We then use the multitask model OpenLAM_2.2.0_27heads_beta3.pt to test the H2O_H2O-PD valid dataset, and it's worth noting that when we use the multitask model to do the testing we need to specify the task head.

• --head: Task head to test if in multi-task mode.

python interface

One can use python interface of DPA-2 to obtain the energy, force, virial of specific structures.

Model Training

We're going to demonstrate how to perform singletask training, multitask training and finetuning based on pretrained model.

singletask

Here we use a small subset of the dataset H2O-PD_DPA_v1_0 to train a singletask model. The file input.json is the input script.

The parameters in input.json are the same as in the previous version of DeePMD-kit except for the descriptor part. For more details, you can go to our paper.

The training can be invoked by

During the training, the error of the model is tested every disp_freq training steps. The training error and validation error are printed correspondingly in the file disp_file (default is lcurve.out). The batch size can be set in the input script by the key batch_size in the corresponding sections for the training and validation data set. An example of the output:

The file contains 8 columns, from left to right, which are the training step, the validation loss, training loss, root mean square (RMS) validation error of energy, RMS training error of energy, RMS validation error of force, RMS training error of force and the learning rate. The RMS error (RMSE) of the energy is normalized by the number of atoms in the system. One can visualize this file with a simple Python script:

multitask

Training on multiple data sets (each data set contains several data systems) can be performed in multi-task mode, with one common descriptor and multiple specific fitting nets for each data set. One needs to switch some parameters in training input script to perform multi-task mode including:

• model –> model_dict, each key of which can be one individual fitting net.

• training_data, validation_data –> data_dict, each key of which can be one individual data set contains several data systems for corresponding fitting net, the keys must be consistent with those in model_dict.

• loss –> loss_dict, each key of which can be one individual loss setting for corresponding fitting net, the keys must be consistent with those in model_dict.

• model_prob, each key of which can be a non-negative integer or float, deciding the chosen probability for corresponding fitting net in training.

Here we use three different datasets(a small subset of the dataset FerroEle_DPA_v1_0, a small subset of the dataset H2O-PD_DPA_v1_0, and a small subset of the dataset SemiCond_DPA_v1_0) to train a multitask model with three task heads.

The training procedure will automatically choose single-task or multi-task mode, based on the above parameters. The training can be invoked by

finetune

Pretraining-and-finetuning is a widely used approach in other fields such as Computer Vision (CV) or Natural Language Processing (NLP) to vastly reduce the training cost, while it’s not trivial in potential models. Compositions and configurations of data samples or even computational parameters in upstream software (such as VASP) may be different between the pretrained and target datasets, leading to energy shifts or other diversities of training data.

The multitask training mode can overcome above difficulties. Our DPA-2 model can hopefully learn the common knowledge in the pretrained dataset and thus reduce the computational cost in downstream training tasks.

Here we have a pretrained multitask model multitask_model.pt on a large dataset (eighteen different datasets), a finetuning strategy can be performed by simply running:

The finetune procedure will inherit the neural network parameters of descriptor in pretrained multitask model. The fitting net can either reinit or inherit the fitting net from any branch of the pre-trained model depending on the argument -m.

• -m (--model-branch): Model branch chosen for fine-tuning if multi-task. If not specified, it will re-init the fitting net.

Whether singletask mode, multitask mode or finetune mode, the training set contains H2O-PD, so we can compare the validation error on dataset H2O-PD directly using a python script

dp freeze

The .pth extension file for molecular dynamics simulations can be obtained by dp freeze.

Molecular Dynamics

The model can drive molecular dynamics in LAMMPS.

Here data.dpa2 gives the initial configuration of water MD simulation, and the file dpa2_in.lammps is the LAMMPS input script. One may check dpa2_in.lammps and finds that it is a rather standard LAMMPS input file for a MD simulation, with only two exception lines:

where the pair style deepmd is invoked and the model file dpa2_model.pth is provided, which means the atomic interaction will be computed by the DPA-2 model that is stored in the file dpa2_model.pth.

In an environment with a compatible version of LAMMPS, the deep potential molecular dynamics can be performed via

Distillation

Distillation can significantly improve the efficiency of finetuned models in MD simulations for production. Distillation requires DP-Gen2. For detail, you can refer to notebook https://bohrium.dp.tech/notebooks/62585747598 or https://bohrium.dp.tech/notebooks/76262686918

DP-Gen based on a DPA-2 pretrained model

Finetuning based on DPA-2 pretrained model can reduce the amount of data required for training. Running DP-Gen with a DPA-2 pretrained model can also save first-principles labelling. DP-Gen with DPA-2 requires DP-Gen2. For detail, you can refer to notebook https://bohrium.dp.tech/notebooks/62585747598 or https://bohrium.dp.tech/notebooks/76262686918

Tips

1. Users are welcome to explore the DP Combo web server , which helps users automate operations such as model training and model distillation. Related notebook: DP Combo教程, 借助DP Combo一键丝滑生成半导体势函数 and 固态电解质实战 | DP Combo@APP体验

2. Current DPA-2 model does not yet support features such as zbl, which we will implement in the near future. If you want to use these features, you can use the previous version of DeePMD-kit(github). Related notebook: DeePMD 使用教程、科研案例、问题收集合集