DOI: 10.5281/zenodo.13342300
title: DPA-2: A Large Atomic Model As a Multi-task Learner

description: Data:





The complete collection of datasets employed in this research is encapsulated within the archive file data-v1.3.tgz. This encompasses both the upstream datasets for pre-training and downstream datasets for fine-tuning, all in DeePMD format. We recommend creating a new directory and employing the command 'tar -xzvf data-v1.3.tgz' to extract the data files.

Inside each dataset contained in subdirectories (e.g., Domains, Metals, H2O, and Others), one will find:



A README file

A 'train' directory (included if utilized in upstream pre-training)



train.json -- A list of file paths for training systems

test.json -- A list of file paths for testing systems



A 'downstream' directory (included if utilized in downstream fine-tuning)



train.json -- A list of file paths for training systems

test.json -- A list of file paths for testing systems



*Main data files comprising various structures

*Additional processing scripts





The root directory contains train.json and downstream.json files that amalgamate the respective upstream and downstream splits mentioned above.




The datasets used in this study are described in Section S1 of the Supplementary Materials and are readily accessible on AIS Square, which provides extensive details.




 

Code:





The 'code' directory, extractable from the archive Code_model_script.tgz, includes the DeePMD-kit's source code, which is based on PyTorch (2.0) Version. Installation and usage instructions can be found within the README file located in deepmd-pytorch-devel.zip.


UPDATE: deepmd-pytorch-devel-0110.zip supports unsupervised learning through denoising, see its README for more details.


 

Model:



Within the 'model' directory, also found in the extracted Code_model_script.tgz, resides the multi-task pre-trained DPA-2 model utilized in this research. Accompanying the model is its configuration file, input.json, which details the simultaneous pre-training of this model across 18 upstream datasets with shared descriptor parameters for 1 million steps.


 

Scripts:



The 'scripts' directory, part of the uncompressed Code_model_script.tgz, comprises all the scripts used for training, fine-tuning (learning curve analysis), and distillation in this work:



1. Upstream_single_task_training: Contains individual training scripts for DPA-2, Gemnet-OC, Equiformer-V2, Nequip, and Allegro, corresponding to the 18 upstream datasets. The training script for MACE is also contained and is applicable across all datasets.

2. Downstream_lcurve_workflow: Includes code and input files to evaluate the learning curves, including tests for DPA-2 fine-tuning transferability across 15 downstream datasets, as depicted in Figure 3 of the manuscript. 

3. Distillation_workflow: Provides input files for distilling the fine-tuned DPA-2 models in datasets such as H2O-PBE0TS-MD, SSE-PBE-D, and FerroEle-D, as illustrated in Figure 4 of the manuscript.



It is important to note that the scripts in 'Upstream_single_task_training' require the installation of deepmd-pytorch and other related models from their respective repositories (Gemnet-OC and Equiformer-V2: here [commit hash: 9bc9373], Nequip: here [commit hash: dceaf49, tag: v0.5.6], Allegro: here [commit hash: 22f673c]), MACE here [commit hash: b76a2a9].

The scripts in 'Downstream_lcurve_workflow' and 'Distillation_workflow' leverage Dflow—a Python framework for constructing scientific computing workflow—and dpgen2, the 2nd generation of the Deep Potential GENerator, both of which are repositories in the Deep Modeling Community.