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. First, install the latest version of DP-Gen2.

This example provides finetuned model and the training and validation data used for finetuning in the dataset. Link them into the working directory.

Then we prepare the initial data for the DP training, i.e. use the finetuned model to label on some data, e.g. the training data used for finetuning.

Then we prepare the initial configurations for MD exploration. Here we sample 100 configurations randomly from the training data.

Below we will prepare the input file for DP-Gen2. You can specify a name for the workflow in the field name. By default, the workflow server https://workflows.deepmodeling.com is used. In bohrium_config, fill in your Bohrium username, password, and project ID. The type_map in the inputs field determines the type map for the final distilled model, and mass_map for the corresponding masses. Specify init_data_sys with the list of system paths for the initial training data we just prepared. valid_data_sys is optional which can be the system paths for the validation data. The training and exploration sections each require input file templates for DP and LAMMPS, which will be provided later. In explore, configurations should be passed with the initial configuration files we just prepared. stages specifies the settings for MD simulations, n_sample determines how many configurations to sample from the initial configurations per iteration, and revisions specifies the values of the variables in the LAMMPS input file template. Each variable's value can be a list, and the final combinations are the Cartesian product of all lists. For more usage of parameters, please refer to the documentation at https://docs.deepmodeling.com/projects/dpgen2/en/latest/.

Here is a simple LAMMPS input template for NVT simulations, where the number of steps, temperature, and output frequency are provided as variables.

This is a DP training input template for distilled model (DPA-1 without attention layer)

Finally, submit the distillation workflow

The progress of the workflow can be tracked through the link printed above. The metrics for each iteration of distillation can be obtained through the dpgen2 command line

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. First, install the latest version of DP-Gen2.

This example provides pretrained model, the initial training data of water and PBE functional files required for VASP calculations in the dataset. Link them into the working directory.

First, we need initial training data. Then we prepare the initial configurations for MD exploration. Here we sample 100 configurations randomly from the training data.

Below we will prepare the input file for DP-Gen2. You can specify a name for the workflow in the field name. By default, the workflow server https://workflows.deepmodeling.com is used. In bohrium_config, fill in your Bohrium username, password, and project ID. Specify init_data_sys with the list of system paths for the initial training data. The training, exploration and first-principle sections each require input file templates for DP, LAMMPS and VASP, which will be provided later. In train, paths of the pretrained models are required for init_models_paths. Here we provide 4 identical paths. In explore, configurations should be passed with the initial configuration files we just prepared. stages specifies the settings for MD simulations, n_sample determines how many configurations to sample from the initial configurations per iteration, and revisions specifies the values of the variables in the LAMMPS input file template. Each variable's value can be a list, and the final combinations are the Cartesian product of all lists. For more usage of parameters, please refer to the documentation at https://docs.deepmodeling.com/projects/dpgen2/en/latest/.

Here is a simple LAMMPS input template for NVT simulations, where the number of steps, temperature, and output frequency are provided as variables.