DPA-3.3-1M

Description

DPA-3.3-1M is trained using a multitask strategy on the OpenLAM datasets V2 with modification. Compared with DPA-3.2-5M, the following changes were made:

1. The model depth was decreased from 24 layers to 6 layers for efficiency.
2. Two new datasets were added to improve model performance in the domain of catalysis. Namely, the CLAM dataset and the OC20mAds dataset.
3. The parameterization of spin multiplicity and charge were decoupled from the fitting net and used as explicit inputs during embedding construction.

Benchmark Results

LAMBench

Note: this tests were carried out with the latest version of LAMBench, new tasks and test datasets were added (code will be released soon). The results may differ from the published version.

Wiggle150

TorsionNet-500

PhononMDR

OC20NEB

How to use

Installation

To use the pretrained model, you need to first install the corresponding version of DeePMD-kit. You can try easy installation:

pip install torch torchvision torchaudio
pip install git+https://github.com/deepmodeling/[email protected]

For other installation options, please refer to the official documentation for off-line installation instructions.

Direct Inference

The pretrained model can be used directly for prediction tasks, such as serving as an ASE calculator.

This model is pretrained in a multi-task manner, featuring a unified backbone (referred to as the unified descriptor and the unified fitting net) and an dataset encoding vector for different tasks. For detailed information, refer to the DPA-3 paper.

The first step involves selecting a specific dataset id from the model to make predictions. To list the available fitting nets (model-branch) in this pretrained model, use the following command:

dp --pt show DPA-3.3-1M.pt model-branch

This will generate an output similar to the following:

Available model branches are ['Domains_Alloy', 'Domains_Anode', 'Domains_Cluster', 'Domains_FerroEle', 
'Domains_SSE_PBE', 'Domains_SemiCond', 'H2O_H2O_PD', 'Metals_AlMgCu', 'Metals_AgAu_PBED3', 'Others_In2Se3',
 'MPGen_OpenCSP', 'Alloy_APEX', 'SSE_ABACUS', 'Hybrid_Perovskite', 'Electrolyte', 'ODAC23', 'Alex2D', 'Omat24', 
'OC20M', 'OC22', 'Organic_Reactions', 'OMol25', 'MPTrj', 'RANDOM'], where 'RANDOM' means using a randomly 
initialized fitting net.

+-------------------+---------------------------+--------------------------------+--------------------------------+
| Model Branch      | Alias                     | description                    | observed_type                  |
+-------------------+---------------------------+--------------------------------+--------------------------------+
| OMat24            | Default, Materials,       | OMat24 is a large-scale open   | ['H', 'He', 'Li', 'Be', 'B',   |
|                   | Omat24, materials, omat24 | dataset containing over 110    | 'C', 'N', 'O', 'F', 'Ne',      |
|                   |                           | million DFT calculations       | 'Na', 'Mg', 'Al', 'Si', 'P',   |
|                   |                           | spanning diverse structures    | 'S', 'Cl', 'Ar', 'K', 'Ca',    |
|                   |                           | and compositions. It is        | 'Sc', 'Ti', 'V', 'Cr', 'Mn',   |
|                   |                           | designed to support AI-driven  | 'Fe', 'Co', 'Ni', 'Cu', 'Zn',  |
|                   |                           | materials discovery by         | 'Ga', 'Ge', 'As', 'Se', 'Br',  |
|                   |                           | providing broad and deep       | 'Kr', 'Rb', 'Sr', 'Y', 'Zr',   |
|                   |                           | coverage of chemical space.    | 'Nb', 'Mo', 'Tc', 'Ru', 'Rh',  |
|                   |                           |                                | 'Pd', 'Ag', 'Cd', 'In', 'Sn',  |
|                   |                           |                                | 'Sb', 'Te', 'I', 'Xe', 'Cs',   |
|                   |                           |                                | 'Ba', 'La', 'Ce', 'Pr', 'Nd',  |
|                   |                           |                                | 'Pm', 'Sm', 'Eu', 'Gd', 'Tb',  |
|                   |                           |                                | 'Dy', 'Ho', 'Er', 'Tm', 'Yb',  |
|                   |                           |                                | 'Lu', 'Hf', 'Ta', 'W', 'Re',   |
|                   |                           |                                | 'Os', 'Ir', 'Pt', 'Au', 'Hg',  |
|                   |                           |                                | 'Tl', 'Pb', 'Bi', 'Ac', 'Th',  |
|                   |                           |                                | 'Pa', 'U', 'Np', 'Pu']         |
+-------------------+---------------------------+--------------------------------+--------------------------------+

Together with the available branch list, a table containing detailed information of each branch is provided. Select a model branch that closely matches your system. Ensure that the elements in your system are included in the corresponding task if you are conducting direct predictions or simulations. For more information on the pretrained datasets, refer to OpenLAM datasets V2.

*Note: If you are unsure which model branch to use, we recommend starting with Omat24, which performs well across a wide range of systems, including inorganic materials, catalytic systems, and molecular systems. If the accuracy of the Omat24 branch does not meet your needs for catalytic applications, consider trying the OC20M branch. For molecular systems, we recommend the OMol25 branch. For 2D inorganic materials, we recommend the Alex2D branch, and for inorganic materials under high pressure, the MPGen_OpenCSP branch. For alloy systems, we recommend the Alloy_APEX branch. *

Assuming you choose OMol25, you can first freeze the model branch from the multi-task pretrained model:

dp --pt freeze -c DPA-3.3-1M.pt -o frozen_model.pth --model-branch  OMol25

or use the alias like

dp --pt freeze -c DPA-3.3-1M.pt -o frozen_model.pth --model-branch  omol25

Then you can use the following Python code for prediction or optimization:

## Compute potential energy
from ase.build import molecule
from deepmd.calculator import DP as DPCalculator
dp = DPCalculator("frozen_model.pth")

singlet = molecule("CH2_s1A1d")

#To set the total charge and spin multiplicity of a structure,
#update charge_spin as [charge, multiplicity]. The default is [0, 1].
singlet.info.update({"charge_spin": [0, 1]})

singlet.calc = dp
print(singlet.get_potential_energy())
print(singlet.get_forces())

## Run BFGS structure optimization
from ase.optimize import BFGS
dyn = BFGS(singlet)
dyn.run(fmax=1e-6)
print(singlet.get_positions())

Zero-Shot Inference

Given that energy labels calculated via DFT can vary by an arbitrary constant, LAMs are consistently used to predict the energy difference between the label and a dummy model that estimates potential energy solely based on the chemical formula. Specifically, the dummy model is defined as

$\qquad E = \sum_e n_{e} \times \hat{b}_e$;

where $E$ is the predicted energy, $n_{e}$ is the number of atoms of element $e$, and $\hat{b}_e$ is the per-atom energy of element $e$ obtained via least-squares fitting:

$ \qquad \hat{b} = \arg\min \sum_i \left( y_i - \sum_e n_{i,e}, b_e \right)^2$,

with $n_{i,e}$ denoting the number of atoms of element $e$ in sample $i$, and $y_i$ the corresponding DFT energy label. When applying the pretrained model to a new downstream system, it is advisable to adjust the energy bias term to better align with the DFT energy labels of the downstream system. This adjustment can enhance prediction accuracy without necessitating further training of the neural network parameters.

To be specific, given a new downstream system, we use dp --pt change-bias to do zero-shot procedure, which will inherit the neural network parameters of the pretrained multitask model, but updating the energy bias to better align with the downstream system.

dp --pt change-bias DPA-3.3-1M.pt -s <your_system> --model-branch OMat24

Fine-tuning

To fine-tune a pretrained model, you need to ensure that the model configuration (model/descriptor and model/fitting_net) matches exactly with the pretrained model. The input_finetune.json file provided on this page contains the same model configuration as the pretrained model.

To prepare for fine-tuning, replace the model/type_map and training/training(validation)_data sections in input_finetune.json with your specific dataset information, just the same as from scratch. Additionally, adjust the learning rate for fine-tuning, typically a small learning rate such as 1e-4 for learning_rate/start_lr is recommended. One can inherit a pretrained fitting net by using the --model-branch argument.

dp --pt train input_finetune.json --finetune DPA-3.3-1M.pt --model-branch OMat24