# Pretraining of attention-based deep learning potential model for molecular
simulation

Duo Zhang  AI for Science Institute, Beijing 100080, P. R. China
DP Technology, Beijing 100080, P. R. China
Academy for Advanced Interdisciplinary Studies, Peking University, Beijing
100871, P. R. China
Hangrui Bi  AI for Science Institute, Beijing 100080, P. R. China
DP Technology, Beijing 100080, P. R. China
Fu-Zhi Dai  AI for Science Institute, Beijing 100080, P. R. China
Wanrun Jiang  AI for Science Institute, Beijing 100080, P. R. China
Linfeng Zhang<EMAIL_ADDRESS>AI for Science Institute, Beijing
100080, P. R. China
DP Technology, Beijing 100080, P. R. China
Han Wang<EMAIL_ADDRESS>Laboratory of Computational Physics, Institute
of Applied Physics and Computational Mathematics, Beijing 100094, P. R. China
HEDPS, CAPT, College of Engineering, Peking University, Beijing 100871, P. R.
China

###### Abstract

Machine learning assisted modeling of the inter-atomic potential energy
surface (PES) is revolutionizing the field of molecular simulation. With the
accumulation of high-quality electronic structure data, a model that can be
pretrained on all available data and finetuned on downstream tasks with a
small additional effort would bring the field to a new stage. Here we propose
DPA-1, a Deep Potential model with a novel attention mechanism, which is
highly effective for representing the conformation and chemical spaces of
atomic systems and learning the PES. We tested DPA-1 on a number of systems
and observed superior performance compared with existing benchmarks. When
pretrained on large-scale datasets containing 56 elements, DPA-1 can be
successfully applied to various downstream tasks with a great improvement of
sample efficiency. Surprisingly, for different elements, the learned type
embedding parameters form a $spiral$ in the latent space and have a natural
correspondence with their positions on the periodic table, showing interesting
interpretability of the pretrained DPA-1 model.

††preprint: APS/123-QED

##  I Introduction

Reliably representing the inter-atomic potential energy surface (PES) is core
to the study of properties of molecules and materials in computational
physics, chemistry, materials science, biology, etc. While electronic
structure methods typically give accurate and transferable PES, they are
prohibitively expensive for scaling to systems of more than thousands of
atoms. On the other hand, empirical force fields are much more efficient, but
are inherently limited by their accuracy in many applications. By properly
integrating machine learning (ML) methodologies and physical requirements like
extensiveness and symmetries, various methods have emerged to address the
accuracy $v.s.$ efficiency dilemma in the realm of PES modeling [1, 2, 3, 4,
5, 6, 7, 8, 9, 10, 11]. Arguably, a new paradigm is forming: electronic
structure methods are no longer used to generate the driving forces during
molecular dynamics simulations but are used to generate data for training
their alternatives, ML-based PES models.

Despite remarkable achievements of ML-based PES models [12, 13, 14],
challenges still remain. For a domain expert who would like to apply such
methodologies in their applications, a natural first question is on the
efforts needed for obtaining a reliable PES model: Are there ready-to-use PES
models? If not, what would be the amount of training data and time cost
required? Can we take advantage of the ever-increasing publicly-available
training data?

To address these issues, there have been several efforts. On one hand,
general-purpose models for various systems, such as silicon [15], phosphorus
[16], water [17], metals and alloys [18, 19, 20, 21, 22], etc., have been
developed and are directly applicable to relevant studies. However, the range
of applicability of such models is typically limited to small conformation or
chemical space. For example, for alloys, the majority of general-purpose ML
models are developed for systems with at most two element types. On the other
hand, several efficient data generation protocols have been developed [23, 24,
25, 26], of which a representative is DP-GEN [25, 26], a concurrent learning
procedure that iteratively explores the configuration space using models
trained with existing data, and then labels only those configurations with
high uncertainty level. Even with these protocols, the computational effort
needed for complicated systems is still prohibitive. For example, to train a
fairly general-purpose model for the AlMgCu alloy system, 100k density
functional theory (DFT) [27, 28] calculations were ultimately performed,
resulting in the cost of ten million CPU core hours [18].

With the accumulation of high-quality electronic structure data covering
almost all the elements on the periodic table, it is becoming possible to
systematically develop pretraining schemes, which have been widely adopted in
areas like computer vision (CV) [29, 30] and natural language processing (NLP)
[31, 32]. In these schemes, one first trains a unified model on large-scale
datasets and then finetunes it for downstream tasks, expecting that a good
representation can be learned in the first stage, and the amount of supervised
data needed for the second stage will be significantly reduced. Recently, the
pretraining-finetuning idea is applied to organic molecules systems for the
energy and force predictions [33, 34], and to tackle tasks beyond representing
the PES [35, 36, 37]. Unfortunately, most ML-based PES models are premature
for such schemes at scale in materials applications. Taking the widely used
two versions of Deep Potential models [6, 7] as examples, the ML parameters
are element-type-dependent, making it highly inefficient when the training
data containing many elements.

Constant efforts have been devoted to adapt the architecture of the ML-based
PES models for large datasets. Among them, one class of models named
equivariant graph neural networks (GNN) [38] that is built upon convolutions
over atomic graphs of node and edge equivariant representations has shown
promise of training on large datasets. SchNet [5], PaiNN [39], GemNet-OC [40],
DimeNet++ [41], PFP [42], SCN [43], SpinConv [44] and Equiformer/EquiformerV2
[45, 46] are trained on the OC20/OC2M [47] dataset containing about 133M/2M
data frames of 56 elements. These models are benchmarked by the accuracy of
energy, force and stable structure predictions. Very recently, it has been
shown that introducing the attention architecture [45] in a GNN model improves
the performance on the OC20/OC2M dataset [46]. Chen and Ong [48] proposed
M3GNet, which was able to train on a subset of the Materials Project [49] that
contains 187,687 configurations encompassing 89 elements and labeled at the
generalized gradient approximation (GGA) [50] or GGA+U level. Takamota et. al.
[42] introduced the PFP model, which was trained on a dataset composed of
molecular and crystal configurations including approximately $9\times 10^{6}$
frames of 45 elements. Choudhary et. al. [51] developed the ALIGNN model, and
they were able to train the model on a subset of the JARVIS-DFT dataset [52]
that is composed of 307,113 data frames of 89 elements. The M3GNet, PFP and
ALIGNN models are proposed as “universal” potential models, however, their
accuracies are not on-par with PES models trained for a specific materials
applications.

The equivariant GNN models are potential candidates for pretraining, several
issues worth special attention before applying them in down-stream real-world
applications. First, the GNN approaches are not well-suited for massively
parallel molecular dynamics simulations [53]. The update of each GNN layer
requires communications between spatially decomposed sub-regions of the
system. In each evaluation of the energy and forces, in total several to a
dozen of such updates are required, which may lead to a substantial
communication overhead in massively parallel high performance super-computers.
Second, some models, such as PaiNN, GemNet-OC, SCN, Equiformer/EquiformerV2,
directly predict forces using rotationally equivariant networks [39, 54, 40,
45] instead of energy gradients with respect to atomic coordinates. Therefore,
the predicted force is not conservative, which serves as a basic assumption in
guaranteeing the accuracy of molecular simulations [55]. Last but not least,
some models, such as GemNet-OC, SpinConv, M3GNet, ALIGNN are not smooth, i.e.
a sudden energy jump may happen as the positions of atoms infinitesimally
varies. This leads to a non-conserved energy in the Hamiltonian dynamics
simulations, which is used in computing the dynamical properties like
diffusion constant and viscosity.

By far, how much may the downstream materials applications benefit from the ML
models trained on the large-scale datasets are still not clear. To answer the
question, in this article, we propose DPA-1, a Deep Potential model with a
novel attention mechanism. Designed with a local descriptor, this model is
exceptionally well-suited for parallel simulations on large-scale systems
containing millions of atoms [56]. Notably, DPA-1 predicts conservative
forces, ensures smoothness, and demonstrates outstanding efficacy in learning
inter-atomic interactions. Moreover, once pretrained, DPA-1 can significantly
decrease the supplementary efforts needed for subsequent downstream tasks. We
tested DPA-1 on various systems and observed superior performance compared
with existing benchmarks. Then we took AlMgCu alloy systems [18] as an
example, showing that after pretraining with single-element and binary
samples, DPA-1 can save around 90% ternary samples compared with the DeepPot-
SE model [7]. Finally, we pretrained DPA-1 using the OC20 dataset, which
consists of 56 elements, and successfully applied it to various downstream
tasks. We checked the interpretability of the pretrained model by looking into
the learned embedding parameters for different element types, finding that the
56 elements are arranged on a $spiral$ in the latent space, which has a
natural correspondence with their physical properties on the periodic table.
Above all, we believe that DPA-1 and the pretraining scheme will bring the
field of molecular simulation to a new stage.

##  II Method

Consider a system of $N$ atoms, the elemental types are
$\mathcal{A}=\left\\{\alpha_{1},\alpha_{2},...,\alpha_{i},...,\alpha_{N}\right\\}$,
and the atomic coordinates are
$\mathcal{R}=\left\\{\bm{r}_{1},\bm{r}_{2},...,\bm{r}_{i},...,\bm{r}_{N}\right\\}$,
with $\bm{r}_{i}$ being the three Cartesian coordinates of atom $i$. The PES
of the system is denoted by $E$, a function of elemental types and
coordinates, i.e. $E=E(\mathcal{A},\mathcal{R})$. For each atom $i$, consider
its neighbors $\\{j|j\in{\mathcal{N}_{r_{c}}(i)}\\}$, where
$\mathcal{N}_{r_{c}}(i)$ denotes the set of atom indices $j$ such that
$r_{ji}<r_{c}$, with $r_{ji}$ being the Euclidean distance between atoms $i$
and $j$. $E$ is represented as the summation of atomic energies
$\left\\{e_{1},e_{2},...,e_{i},...,e_{N}\right\\}$, where the atomic energy
$e_{i}$ only depends on the information of ${\mathcal{N}_{r_{c}}(i)}$. We
define $N_{i}=|\mathcal{N}_{r_{c}}(i)|$, the cardinality of the set
$\mathcal{N}_{r_{c}}(i)$. We use $\mathcal{A}^{i}$ to denote element types in
${\mathcal{N}_{r_{c}}(i)}$, and $\mathcal{R}^{i}\in\mathbb{R}^{N_{i}\times 3}$
their corresponding coordinates relative to $i$. The atomic energy $e_{i}$ is
thus a function of $\mathcal{A}^{i}$ and $\mathcal{R}^{i}$. The atomic force
on atom $i$, $\mathcal{F}_{i}$, is defined as the negative gradient of the
total energy with respect to $i$’s coordinate:

$\mathcal{F}_{i}=-\nabla_{\bm{r}_{\bm{i}}}E.$ (1)

We refer to Ref. [7] for a detailed discussion on several requirements on PES
modeling. In particular, the PES has to be invariant under translation,
rotation, and permutation of the indices of atoms with the same element types.

The details of the model architecture are introduced below. We refer to Fig. 1
for the overall pipeline to predict the atomic energy $e_{i}$: from the
embedded neighboring environment, through the self-attention scheme, to the
symmetry-preserving descriptors, and finally to the fitting network.

Figure 1: Schematic illustration of DPA-1. (a) Flowchart from
$\mathcal{A}^{i}$ and $\mathcal{R}^{i}$ to the atomic energy $e_{i}$. (b)
Structure of the Embedding net, which maps $s(r_{ji})$ and $T_{i}$, through
multiple residual layers, to $\mathcal{G}^{i}$. (c) Self-attention mechanism
on $\mathcal{G}^{i}$ through a standard scale-dot procedure gated by the
angular information $\hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T}$. (d)
Fitting net structures, similar to Embedding net, from the descriptors
$\mathcal{D}^{i}$ and $T_{i}$ to final atomic energy $e_{i}$.

###  II.1 Local embedding matrix with type information

We obtain the local embedding matrix with the following three steps. First,
$\mathcal{R}^{i}$ is mapped to the generalized coordinates
$\tilde{\mathcal{R}}^{i}\in\mathbb{R}^{N_{i}\times 4}$. In this mapping, each
row of $\mathcal{R}^{i},\left\\{x_{ji},y_{ji},z_{ji}\right\\}$, is transformed
into a row of $\tilde{\mathcal{R}}^{i}$:

$\left\\{x_{ji},y_{ji},z_{ji}\right\\}\mapsto\left\\{s\left(r_{ji}\right),\hat{x}_{ji},\hat{y}_{ji},\hat{z}_{ji}\right\\},$
(2)

where $\left\\{x_{ji},y_{ji},z_{ji}\right\\}$ denotes the Cartesian
coordinates of $\bm{r_{ji}}=\bm{r_{j}}-\bm{r_{i}}$,
$\hat{x}_{ji}=\frac{s\left(r_{ji}\right)x_{ji}}{r_{ji}},\hat{y}_{ji}=\frac{s\left(r_{ji}\right)y_{ji}}{r_{ji}},\hat{z}_{ji}=\frac{s\left(r_{ji}\right)z_{ji}}{r_{ji}}$,
and $s\left(r_{ji}\right):\mathbb{R}\mapsto\mathbb{R}$ is a continuous and
differentiable scalar weighting function applied to each component, defined
as:

$s\left(r_{ji}\right)=\begin{cases}\frac{1}{r_{ji}}&r_{ji}<r_{cs}\\\
\frac{1}{r_{ji}}\left[u^{3}\left(-6u^{2}+15u-10\right)+1\right]&r_{cs}\leq
r_{ji}<r_{c},\quad u=\frac{r_{ji}-r_{cs}}{r_{c}-r_{cs}}.\\\ 0&r_{c}\leq
r_{ji}\end{cases}$ (3)

Here $r_{cs}$ is a smooth cutoff parameter that allows the components in
$\tilde{\mathcal{R}}^{i}$ to smoothly go to zero at the boundary of the local
region defined by $r_{c}$.

Second, we add the atomic type embedding as supplemental information. For atom
$i$, the type embedding map $T_{i}$ is defined as:

$T_{i}=\phi_{T}(\alpha_{i}),$ (4)

where $\alpha_{i}$ is the atomic type of atom $i$ and $\phi_{T}$ is a one-hot-
like embedding network mapping from $\alpha_{i}$ to a length-fixed vector.

Then, given both $\tilde{\mathcal{R}}^{i}$ and type embeddings
$\left\\{T_{i}\right\\}\cup\left\\{T_{j}|j\in{\mathcal{N}_{r_{c}}(i)}\right\\}$,
we define the local embedding matrix
$\mathcal{G}^{i}\in\mathbb{R}^{N_{i}\times M_{1}}$:

$\left(\mathcal{G}^{i}\right)_{j}=G\left(s\left(r_{ji}\right),T_{i},T_{j}\right),$
(5)

where $G$ is a neural network mapping from scalar weight
$s\left(r_{ji}\right)$ and type embeddings of both center and neighbor atoms,
through multiple hidden layers, to $M_{1}$ outputs. Here we simply feed the
concatenated inputs into $G$ at once, as shown in Fig. 1(b).

###  II.2 Attention method for building up trainable descriptors

The attention mechanism has achieved great success and played an increasingly
important role in CV [57] and NLP [58]. It has become an excellent tool for
modeling the importance or relevance of visual regions or text tokens, thus is
potentially appropriate to reweight the interactions among neighbor atoms
according to both distance and angular information.

In DPA-1, we follow the standard self-attention mechanism and obtain the
queries $\mathcal{Q}^{i,l}\in\mathbb{R}^{N_{i}\times d_{k}}$, keys
$\mathcal{K}^{i,l}\in\mathbb{R}^{N_{i}\times d_{k}}$, and values
$\mathcal{V}^{i,l}\in\mathbb{R}^{N_{i}\times d_{v}}$:

$\displaystyle\left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),$
$\displaystyle\left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),$
(6)
$\displaystyle\left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),$

where $Q_{l},K_{l},V_{l}$ represent three linear transformations which output
the queries and keys of dimension $d_{k}$ and values of dimension $d_{v}$, and
$l$ is the index of attention layer. Here we take
$\mathcal{G}^{i,0}=\mathcal{G}^{i}$.

Then we adopt the scaled dot-product attention method[59] to mix the neighbor
features after calculating the attention weights:

$A(\mathcal{Q}^{i,l},\mathcal{K}^{i,l},\mathcal{V}^{i,l},\mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l},\mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l},$
(7)

where
$\varphi\left(\mathcal{Q}^{i,l},\mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\in\mathbb{R}^{N_{i}\times
N_{i}}$ is attention weights. In the original attention method, one typically
has
$\varphi\left(\mathcal{Q}^{i,l},\mathcal{K}^{i,l}\right)=\operatorname{softmax}\left(\frac{\mathcal{Q}^{i,l}(\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)$,
with $\sqrt{d_{k}}$ being the normalization temperature. This is slightly
modified to better incorporate the angular information:

$\varphi\left(\mathcal{Q}^{i,l},\mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)=\operatorname{softmax}\left(\frac{\mathcal{Q}^{i,l}(\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)\odot\hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T},$
(8)

where
$\hat{\mathcal{R}}^{i}=\frac{\mathcal{R}^{i}}{\|\mathcal{R}^{i}\|_{2}}\in\mathbb{R}^{N_{i}\times
3}$ denotes normalized relative coordinates and $\odot$ means element-wise
multiplication. Intuitively, in the neighborhood of center atom $i$, neighbor
atom $k$ may be highly correlated with $j$ when both the relative distance
attention $(\mathcal{Q}^{i,l})_{j}(\mathcal{K}^{i,l})_{k}^{T}$ and normalized
product of relative coordinates
$\frac{\textbf{r}_{ji}(\textbf{r}_{ki})^{T}}{r_{ji}r_{ki}}$ have high scores.

Then we add layer normalization in a residual way to finally obtain the self-
attentioned local embedding matrix $\hat{\mathcal{G}}^{i}$ in one such
attention layer:

$\mathcal{G}^{i,l}=\mathcal{G}^{i,l-1}+\mathrm{LayerNorm}(A(\mathcal{Q}^{i,l},\mathcal{K}^{i,l},\mathcal{V}^{i,l},\mathcal{R}^{i,l})).$
(9)

We also tried other attention related tricks such as pre-layer normalization,
multi-head attention, etc., which brought little improvement. In practice, as
shown in Fig. 1(c), we repeated this procedure by $l(l\geq 2)$ times for a
more complete representation. If not stated otherwise, we use $l=2$ in the
following sections of the work.

Next, we define the encoded feature matrix
$\mathcal{D}^{i}\in\mathbb{R}^{M_{1}\times M_{2}}$ of atom $i$:

$\mathcal{D}^{i}=\left(\hat{\mathcal{G}}^{i}\right)^{T}\tilde{\mathcal{R}}^{i}\left(\tilde{\mathcal{R}}^{i}\right)^{T}\dot{\mathcal{G}}^{i},$
(10)

where $\dot{\mathcal{G}}^{i}$ stands for a sub-matrix of
$\hat{\mathcal{G}}^{i}$, which takes the first $M_{2}(<M_{1})$ columns of
$\hat{\mathcal{G}}^{i}$. Here the feature matrix $\mathcal{D}^{i}$, i.e. the
descriptor, preserves all the invariance mentioned above, of which the proof
can be found in Ref. [7]. We then pass the reshaped $\mathcal{D}^{i}$,
concatenated with the type embedding parameters of the center atom, through
the multi-layer fitting network:

$e_{i}=e\left(\mathcal{D}^{i},T_{i}\right).$ (11)

The total energy of the system is then given as the summation of $e_{i}$, and
the atomic force $\mathcal{F}_{i}$ can be further computed via Eq.1.

###  II.3 Model (pre-)training and finetuning

For model training or pretraining, we adopted the Adam stochastic gradient
descent method [60] on all the trainable parameters $\bm{w}$ inside the model
to minimize the loss:

$\mathcal{L}_{\bm{w}}(E^{\bm{w}},\mathcal{F}^{\bm{w}})=\frac{1}{|\mathcal{B}|}\sum_{t\in\mathcal{B}}\left(p_{\epsilon}\left|E_{t}-E_{t}^{\bm{w}}\right|^{2}+p_{f}\left|\mathcal{F}_{t}-\mathcal{F}_{t}^{\bm{w}}\right|^{2}\right).$
(12)

Here $\mathcal{B}$ represents a minibatch, $|\mathcal{B}|$ is the batch size,
$t$ denotes the index of the training sample.
$E^{\bm{w}},\mathcal{F}^{\bm{w}}$ denote the model outputs and $E,\mathcal{F}$
are the corresponding DFT results. We also adopted a scheduler to tune the
prefactors $p_{\epsilon}$ and $p_{f}$ during the training process to make a
better balance between energy and force labels. Virial errors, which are
omitted here, can be added to the loss for training if available.

To finetune the pretrained model with a new dataset, we first change the
energy bias in the last layer of the pretrained model with the new statistical
results of the new dataset, and then we fix part of the parameters in the
pretrained model and train the remaining. For the following experiments, we
obtained the best performance when only the type embedding parameters are
fixed.

##  III Experiments

We conducted a number of experiments to evaluate the performance of DPA-1.
First, to test the model’s ability to transfer among different compositions,
we trained it from scratch against various systems and tested it under several
challenging schemes. Then, we used an AlMgCu dataset to test its ability to
transfer to ternary systems upon pretraining with single-element and binary
data. Finally, we pretrained DPA-1 using the OC2M subset in OC20 dataset [47]
and applied it to various downstream tasks. To illustrate the effectiveness of
the type-embedding and attention schemes, we compared against DeepPot-SE model
[7] in all the experiments. In the following, we shall introduce first the
datasets we used, and then the experiments we conducted.

Table 1: Validation RMSE of DPA-1 and DeepPot-SE on energy (meV/atom) and atomic forces (meV/Å) with different settings of the training/validation sets (See Sec. III.1 for details). The number of attention layers $l$ in DPA-1 is set to 2 in the AlMgCu and SSE systems, and to 3 in the HEA systems. Bold numbers correspond to lower values.  Systems | Training | Validation | Validation RMSE
---|---|---|---
DPA-1 | DeepPot-SE
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
AlMgCu | single + binary | ternary | 6.99 | 58 | 65.1 | 92
all (single + binary + ternary) | ternary | 2.26 | 35 | 3.16 | 42
all | all | 2.74 | 38 | 3.67 | 45
SSE | init + single | mix | 0.56 | 60 | 0.72 | 76
init + mix | single | 3.72 | 69 | 3.76 | 82
all (init + single + mix) | all | 1.41 | 68 | 2.92 | 85
HEA | interior | exterior | 31.2 | 158 | 197 | 399
exterior | interior | 6.88 | 117 | 236 | 428
all (interior + exterior) | all | 4.96 | 71 | 28.7 | 141
Figure 2:  Learning curves of both energy and force with DeepPot-SE and DPA-1,
under different setups and on different systems. (a) Learning curves on the
AlMgCu ternary subset, with DeepPot-SE and DPA-1 models pretrained on single-
element and binary subsets; (b-c) Learning curves on HEA (b) and AlCu (c),
with DeepPot-SE (from scratch) and DPA-1 (both from scratch and pretrained on
OC2M). Red line represents the full-data-training baseline with DPA-1.  Figure
3:  (a) 3-dimensional PCA visualization of the learned type embeddings of
DPA-1 pretrained on OC2M. These 56 elements are roughly arranged on a _spiral_
in the latent space. Elements in the fourth period are connected with the red
line and elements belonging to the same family are grouped by the blue dot
lines. Colors on the names of the elements represent the height in $z$-axis.
We use dashed circle to denote the hypothetical position of Li, which is not
contained in OC2M. See text for discussions. (b) RMSE of energy and force for
SSE systems given by DPA-1 pretrained on OC2M, as functions of linear
interpolation coefficient $\lambda\left(Na\right)$. Since Li is not contained
in OC2M, we let
$T_{Li}=\lambda\left(Na\right)*T_{Na}+\left(1-\lambda\left(Na\right)\right)*T_{H}$
be the interpolated type embedding of Li. The OC2M-pretrained model with this
interpolation and modified energy bias is directly tested on SSE systems
without further training. (c) Training efficiency of DPA-1 and DeepPot-SE
(considering type information of both two sides) with the growing number of
element types in training systems. The maximum number of neighboring atoms to
be considered is set to 120 in all the experiments.

###  III.1 Datasets

AlMgCu alloy systems [18]. This dataset is generated using DP-GEN [26], a
concurrent learning scheme. After exploring 2.73 billion alloy configurations
(derived from $\sim$2000 bulk and surface systems), only a small portion
($\sim$100k configurations ) of them are labeled then compose the compact
dataset. The exploration runs in the whole concentration space, i.e.,
Al_xCu_yMg_z with $0\leq x,y,z\leq 1,x+y+z=1$, and $x$, $y$, $z$ take discrete
values permitted by the finite-size simulation boxes. We can divide the
systems into _single_ , _binary_ and _ternary_ subsets, in name of the number
of non-zero $x$, $y$, $z$. The configuration space covers a temperature range
around 50.0 K to 2579.8 K and a pressure range around 1 bar to 50000 bar.

Solid-state electrolyte (SSE) systems [61]. These systems contain
Li10XP2S12-type SSE materials, where X represents single or combination of
Ge/Si/Sn, and can be divided into three main parts: _init_ , _mix_ and
_single_. The _init_ part comes from standard DP-GEN scheme starting from 590
structures that are generated via slightly perturbing DFT-relaxed crystal
structures Li10Ge(PS6)2, Li10SiP2S12 and Li10SnP2S12 from Materials Project
[49]. The exploration covers both ordered structures relaxed by DFT (i.e.
structures downloaded from the Materials Project database, in which the
position of Ge/Si/Sn/P atoms are fixed) and disordered structures whose 4d
sites are randomly occupied by Ge/Si/Sn/P. Based on the _init_ part, the _mix_
part contains further exploration in binary and ternary mixture of Ge/Si/Sn,
while the _single_ part covers only single X in Ge/Si/Sn with other changes in
lattice and ratio of Li.

HEA systems. The High Entropy Alloy HEA dataset includes bulk TaNbWMoVAl alloy
systems of various configurations and compositions. We employ DP-GEN to
explore the composition space, starting from Ta_3Nb_3W_3Mo_3V_3Al_1, a 16-atom
unit cell containing the former 5 elements as main components and Al as an
additive. The dataset is divided into two subsets: _interior_ and _exterior_.
The _interior_ (higher entropy) subset includes composition variations near
the starting point. It covers six-component, quinary, quaternary and ternary
alloys. The _exterior_ (lower entropy) subset includes systems that are close
to the corners and edges of the composition space. It includes systems where
one or two elements dominate, binary alloys and simple substance systems. For
both subsets, the temperature range is around 50.0 K to 388.1 K and the
pressure range is around 1 bar to 50000 bar.

OC20 [47]. OC20 consists of single adsorbates (small molecules) physically
binding to the surfaces of catalysts covering periodic bulk materials with 56
elements. Both the chemical diversity and system size are much more complex
than other benchmark datasets, such as MD17 [62], ANI-1x [24] or QM9 [63].
OC2M is a subset including 2 million data points (energies and forces)
randomly sampled from OC20, which is still challenging for model training and
decent for pretraining. Johannes et al. recently provided several baselines on
OC2M, taking months to converge[40].

###  III.2 Accuracy on various datasets, trained from scratch

The majority of existing models usually focus on the ability to transfer among
different configurations, in which case training and validation subsets
consist of similar compositions (e.g. randomly sampled from the same dataset).
However, to perform pretraining, the upstream and downstream datasets may
differ violently. Thus, it’s vital for models under pretraining scheme to
transfer among different compositions or even among different datasets, which
has, as far as we know, rarely been discussed before. In this work, we mainly
focus on a more general but challenging scheme to comprehensively test the
generalization ability of the model.

We first designed several challenging tasks to test the model’s ability to
transfer among different compositions. For AlMgCu, SSE, and HEA systems, we
divided them into subsets with different compositions for training and
validation (See Sec. III.1 for details). The results of DPA-1 and DeepPot-SE
are shown in Table 1. With the training loss nearly the same (omitted in the
table), the DPA-1 drastically outperforms DeepPot-SE in the validation
accuracy. For example, for AlMgCu systems, when trained only on single- and
binary-element samples, the validation RMSE of DPA-1 on ternary samples can
outperform DeepPot-SE by one order of magnitude (6.99 versus 65.1 meV/atom).
This suggests that the DPA-1 model might have learned the latent interactions
of ternary pairs Al-Mg-Cu from binary pairs Al-Mg, Al-Cu, Mg-Cu, and single-
element interactions, possibly thanks to the type-embedding scheme and
attention mechanism. We conducted an ablation study in Appendix. A on HEA
systems to demonstrate the influence of each structural component.

To test the performance of DPA-1 in terms of predicting more physical
quantities, we performed geometry relaxations on all AlMgCu ternary alloys
available from the Materials Project to evaluate their accuracy in predicting
formation energy and equilibrium volume (see details in Appendix B). We also
used it to calculate the elastic moduli of AlMgCu systems, which requires
accurately capturing the second-order information (see details in Appendix C).
Additionally, we carried out molecular dynamics simulations on LiGePS systems
to assess the diffusion coefficients in relation to temperature, comparing the
results to ab initio molecular dynamics (AIMD) simulations and experimental
studies (see details in Appendix D). In all tests, satisfactory agreement with
the DFT and/or experimental references are obtained.

As a supplement, we also trained DPA-1 model on several simple systems to
compare with other ML-based PES. Since these tasks are much easier than the
above ones and out of our main focus, we place the results in Appendix Section
F. Note that there may be relatively little room for improvement on these
simple datasets, while DPA-1 can still outperform other methods with even less
training samples.

###  III.3 Sample efficiency of pretrained models

As shown in Fig. 2, we use the learning curves to illustrate in terms of the
amount of additional training data saved for downstream tasks thanks to model
pretraining. In all the experiments, the learning curves were generated by an
active learning procedure, in which a pool of data labeled by energy and force
is prepared and three steps are repeated iteratively: using samples in the
training pool to train the model; testing the model using the remaining
samples; selecting 50 samples with the largest prediction errors on per-atom
energies and adding them to the training pool. We use the term sample
efficiency to denote the amount of training samples required by a model to
achieve a given accuracy level for a certain task.

We started with a relatively simple task to compare DeepPot-SE and DPA-1. In
this task, both the two models were pretrained using single-element and binary
subsets of the AlMgCu systems, and the learning curves were obtained using the
AlMgCu ternary subset. As shown in Fig. 2(a), DPA-1 exhibits a much better
sample efficiency than DeepPot-SE, which should be expected.

Next, we used the OC2M dataset, which contains 56 elements, to pretrain DPA-1
and evaluated its performance on the HEA systems and the AlCu systems (Figs.
2(b) and (c), respectively). As shown in Fig. 3(c), the training cost of
DeepPot-SE scales quadratically with the number of elements, making its
pretraining computationally infeasible, while the number of elements has no
effects on the training cost of DPA-1. It is observed that the sample
efficiency of DPA-1 pretrained on OC2M is generally better than DPA-1 from
scratch, while DeepPot-SE from scratch is the worst. Moreover, compared with
AlCu systems, the improvement of pretraining is much more significant for HEA
systems, possibly due to the fact that the number of elements of HEA is much
larger than AlCu, and the local chemical environment is much more complicated.

The equivariant GNN models usually needs thousands of GPU hours to be trained
to a descent accuracy [40]. By contrast, the DPA-1 modle only takes less than
200 GPU hours for training. The converged energy and force MAEs on the OC2M
validation set are 0.681 eV and 0.076 eV/Å, respectively. This accuracy is
comparable with the best energy conserving GNN model DimeNet++, which achieves
MAEs of 0.805 eV and 0.066 eV/Å, reported in Ref. [40]. A better performance
of energy MAE 0.286 eV and force MAE 0.026 eV/Å is achieved by GemNet-OC at
the cost of non-conservative forces and loss of smoothness [40].

###  III.4 Interpretability of type embedding learned from pretraining

To see whether DPA-1 can learn physically meaningful information from
pretraining, we investigated the 3-dimensional principal component analysis
(PCA) visualization of the learned type embeddings in the OC2M-pretrained
model. Interestingly, as shown in Fig. 3(a), the arrangement of the elements
generally follows the shape of a downward spiral. Elements belonging to the
same period are lined up in the direction of the spiral; while elements
belonging to the same family are listed in the direction orthogonal to the
spiral. Even though some transition metal elements are almost bounded
together, this rule still roughly holds. It is observed that C, N and O are
outliers, possibly because in OC2M, C, N and O are mostly in organic
molecules, which serve as adsorbates and have chemical environments that are
very different from other elements.

In addition, we performed interpolation experiments for the type embedding of
Li, an element unseen in OC2M. As shown in Fig. 3(b), we let
$T_{Li}=\lambda\left(Na\right)*T_{Na}+\left(1-\lambda\left(Na\right)\right)*T_{H}$,
since Li lies between H and Na in the same family. When tested on the SSE
system, only the bias in the atomic energy is changed, since the setup of the
electronic method used to label the SSE system is different from that for
OC2M, which typically causes an energy shift. It is found that the RMSE of
energy and force shows a sudden drop when $\lambda\left(Na\right)=0.7$, which
meets the chemical intuition and further confirms the interpretability of the
pretrained DPA-1 model. Moreover, we conducted analogous interpolation
experiments for Nb and Mo on the HEA systems, and reached similar conclusions
as the Li interpolation (see detailed report in Appendix. E).

##  IV Summary

In this paper, we developed DPA-1, an attention-based Deep Potential model
that allows for large-scale pretraining on atomistic datasets. We tested DPA-1
from different aspects, showing its excellent performance in terms of its
accuracy on various datasets when trained from scratch, as well as its sample
efficiency when pretrained with existing data. Further investigations on the
type embedding parameters suggests the interpretability of DPA-1 pretrained on
OC2M.

In the future, it will be of interest to extend the training dataset to cover
the full periodic table, and, in particular, see a more converged “spiral” in
the latent space; the embedding information of local chemical environments may
be useful to characterize different conformations. Multi-task and unsupervised
training schemes would worth exploring; and, for downstream tasks, just like
what has happened in the fields of CV and NLP, schemes like model compression,
distillation, and transfer, etc., are desperately needed. We leave these
possibilities and more applications to future works.

##  V Data availability

The dataset used for training OC2M-pretrained DPA-1 is available at:
https://www.aissquare.com/datasets/detail?pageType=datasets&name=Open_Catalyst_2020(OC20_Dataset).
Other datasets are available in their references or on reasonable request.

##  VI Code availability

The codes of DPA-1 are in the repository of DeePMD-kit:
https://github.com/deepmodeling/deepmd-kit. The OC2M-pretrained model is
available at:
https://www.aissquare.com/models/detail?pageType=models&name=DPA_1_OC2M. All
the experiments are performed on single Nvidia Tesla V100.

##  VII AUTHOR CONTRIBUTIONS

D.Z, L.Z, H.W. and F.Z.D. conceived the idea of this work. D.Z., H.B. and H.W.
designed the model structure. D.Z. implemented the model. D.Z., H.B. and W.J
performed the experiments on different systems. All authors contributed to the
discussions and edited the manuscript.

##  VIII COMPETING INTERESTS

The authors declare no competing financial or non-financial interests.

###### Acknowledgements.

The work of H.W. was supported by the National Science Foundation of China
under Grant No.11871110 and 12122103. We thank Y.L., Z.L. and G.K. for
inspiring discussions. The computational resource was supported by the Bohrium
Cloud Platform at DP technology.

## References

  * [1] Jörg Behler and Michele Parrinello.  Generalized neural-network representation of high-dimensional potential-energy surfaces.  Physical review letters, 98(14):146401, 2007.
  * [2] Albert P Bartók, Mike C Payne, Risi Kondor, and Gábor Csányi.  Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons.  Physical review letters, 104(13):136403, 2010.
  * [3] Aidan P Thompson, Laura P Swiler, Christian R Trott, Stephen M Foiles, and Garritt J Tucker.  Spectral neighbor analysis method for automated generation of quantum-accurate interatomic potentials.  Journal of Computational Physics, 285:316–330, 2015.
  * [4] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  In International conference on machine learning, pages 1263–1272. PMLR, 2017.
  * [5] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller.  Schnet: A continuous-filter convolutional neural network for modeling quantum interactions.  Advances in neural information processing systems, 30, 2017.
  * [6] Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and EJPRL Weinan.  Deep potential molecular dynamics: a scalable model with the accuracy of quantum mechanics.  Physical review letters, 120(14):143001, 2018.
  * [7] Linfeng Zhang, Jiequn Han, Han Wang, Wissam Saidi, Roberto Car, et al.  End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems.  Advances in Neural Information Processing Systems, 31, 2018.
  * [8] Ralf Drautz.  Atomic cluster expansion for accurate and transferable interatomic potentials.  Physical Review B, 99(1):014104, 2019.
  * [9] Johannes Gasteiger, Janek Groß, and Stephan Günnemann.  Directional message passing for molecular graphs.  In International Conference on Learning Representations, 2019.
  * [10] Yaolong Zhang, Ce Hu, and Bin Jiang.  Embedded atom neural network potentials: Efficient and accurate machine learning with a physically inspired representation.  The journal of physical chemistry letters, 10(17):4962–4967, 2019\.
  * [11] Johannes Gasteiger, Florian Becker, and Stephan Günnemann.  Gemnet: Universal directional graph neural networks for molecules.  Advances in Neural Information Processing Systems, 34:6790–6802, 2021.
  * [12] Volker L Deringer, Albert P Bartók, Noam Bernstein, David M Wilkins, Michele Ceriotti, and Gábor Csányi.  Gaussian process regression for materials and molecules.  Chemical Reviews, 121(16):10073–10141, 2021.
  * [13] Oliver T Unke, Stefan Chmiela, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schütt, Alexandre Tkatchenko, and Klaus-Robert Müller.  Machine learning force fields.  Chemical Reviews, 121(16):10142–10186, 2021.
  * [14] Tongqi Wen, Linfeng Zhang, Han Wang, Weinan E, and David J Srolovitz.  Deep potentials for materials science.  Materials Futures, 2022.
  * [15] Albert P. Bartók, James Kermode, Noam Bernstein, and Gábor Csányi.  Machine learning a general-purpose interatomic potential for silicon.  Phys. Rev. X, 8:041048, Dec 2018.
  * [16] Volker L Deringer, Miguel A Caro, and Gábor Csányi.  A general-purpose machine-learning force field for bulk and nanostructured phosphorus.  Nature communications, 11(1):1–11, 2020.
  * [17] Linfeng Zhang, Han Wang, Roberto Car, and Weinan E.  Phase diagram of a deep potential water model.  Phys. Rev. Lett., 126:236001, Jun 2021.
  * [18] Wanrun Jiang, Yuzhi Zhang, Linfeng Zhang, and Han Wang.  Accurate deep potential model for the al–cu–mg alloy in the full concentration space.  Chinese Physics B, 30(5):050706, 2021.
  * [19] Wojciech J. Szlachta, Albert P. Bartók, and Gábor Csányi.  Accuracy and transferability of gaussian approximation potential models for tungsten.  Phys. Rev. B, 90:104108, Sep 2014.
  * [20] Xiaoyang Wang, Yinan Wang, Linfeng Zhang, Fuzhi Dai, and Han Wang.  A tungsten deep neural-network potential for simulating mechanical property degradation under fusion service environment.  Nuclear Fusion, 2022.
  * [21] YiNan Wang, LinFeng Zhang, Ben Xu, XiaoYang Wang, and Han Wang.  A generalizable machine learning potential of Ag–Au nanoalloys and its application to surface reconstruction, segregation and diffusion.  Modelling and Simulation in Materials Science and Engineering, 30(2):025003, dec 2021.
  * [22] Tongqi Wen, Rui Wang, Linyu Zhu, Linfeng Zhang, Han Wang, David J. Srolovitz, and Zhaoxuan Wu.  Specialising neural network potentials for accurate properties and application to the mechanical response of titanium.  npj computational materials, 7:206, 2021.
  * [23] Evgeny V Podryabinkin and Alexander V Shapeev.  Active learning of linearly parametrized interatomic potentials.  Computational Materials Science, 140:171–180, 2017.
  * [24] Justin S Smith, Ben Nebgen, Nicholas Lubbers, Olexandr Isayev, and Adrian E Roitberg.  Less is more: Sampling chemical space with active learning.  The Journal of chemical physics, 148(24):241733, 2018.
  * [25] Linfeng Zhang, De-Ye Lin, Han Wang, Roberto Car, and E Weinan.  Active learning of uniformly accurate interatomic potentials for materials simulation.  Physical Review Materials, 3(2):023804, 2019.
  * [26] Yuzhi Zhang, Haidi Wang, Weijie Chen, Jinzhe Zeng, Linfeng Zhang, Han Wang, and E Weinan.  Dp-gen: A concurrent learning platform for the generation of reliable deep learning based potential energy models.  Computer Physics Communications, 253:107206, 2020.
  * [27] Walter Kohn and Lu Jeu Sham.  Self-consistent equations including exchange and correlation effects.  Physical review, 140(4A):A1133, 1965.
  * [28] Roberto Car and Michele Parrinello.  Unified approach for molecular dynamics and density-functional theory.  Physical review letters, 55(22):2471, 1985.
  * [29] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  International journal of computer vision, 115(3):211–252, 2015\.
  * [30] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [31] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [32] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  Advances in neural information processing systems, 33:1877–1901, 2020.
  * [33] Justin S Smith, Olexandr Isayev, and Adrian E Roitberg.  Ani-1: an extensible neural network potential with dft accuracy at force field computational cost.  Chemical science, 8(4):3192–3203, 2017.
  * [34] Justin S Smith, Benjamin T Nebgen, Roman Zubatyuk, Nicholas Lubbers, Christian Devereux, Kipton Barros, Sergei Tretiak, Olexandr Isayev, and Adrian E Roitberg.  Approaching coupled cluster accuracy with a general-purpose neural network potential through transfer learning.  Nature communications, 10(1):1–8, 2019.
  * [35] Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang.  Pre-training molecular graph representation with 3d geometry.  In International Conference on Learning Representations, 2022.
  * [36] Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò.  3d infomax improves gnns for molecular property prediction.  arXiv preprint arXiv:2110.04126, 2021.
  * [37] Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke.  Uni-mol: A universal 3d molecular representation learning framework.  2022\.
  * [38] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley.  Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds.  arXiv preprint arXiv:1802.08219, 2018.
  * [39] Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger.  Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021.
  * [40] Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zitnick, and Abhishek Das.  Gemnet-oc: Developing graph neural networks for large and diverse molecular simulation datasets, 2022.
  * [41] Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann.  Fast and uncertainty-aware directional message passing for non-equilibrium molecules, 2022.
  * [42] So Takamoto, Chikashi Shinagawa, Daisuke Motoki, Kosuke Nakago, Wenwen Li, Iori Kurata, Taku Watanabe, Yoshihiro Yayama, Hiroki Iriguchi, Yusuke Asano, et al.  Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements.  Nature Communications, 13(1):2991, 2022.
  * [43] C. Lawrence Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood.  Spherical channels for modeling atomic interactions, 2022.
  * [44] Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitnick.  Rotation invariant graph neural networks using spin convolutions, 2021\.
  * [45] Yi-Lun Liao and Tess Smidt.  Equiformer: Equivariant graph attention transformer for 3d atomistic graphs, 2023.
  * [46] Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt.  Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations, 2023.
  * [47] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, et al.  Open catalyst 2020 (oc20) dataset and community challenges.  ACS Catalysis, 11(10):6059–6072, 2021.
  * [48] Chi Chen and Shyue Ping Ong.  A universal graph deep learning interatomic potential for the periodic table.  arXiv preprint arXiv:2202.02450, 2022.
  * [49] Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, et al.  Commentary: The materials project: A materials genome approach to accelerating materials innovation.  APL materials, 1(1):011002, 2013.
  * [50] John P Perdew, Kieron Burke, and Matthias Ernzerhof.  Generalized gradient approximation made simple.  Physical review letters, 77(18):3865, 1996.
  * [51] Kamal Choudhary, Brian DeCost, Lily Major, Keith Butler, Jeyan Thiyagalingam, and Francesca Tavazza.  Unified graph neural network force-field for the periodic table: solid state applications.  Digital Discovery, 2(2):346–355, 2023.
  * [52] Kamal Choudhary, Kevin F Garrity, Andrew CE Reid, Brian DeCost, Adam J Biacchi, Angela R Hight Walker, Zachary Trautt, Jason Hattrick-Simpers, A Gilad Kusne, Andrea Centrone, et al.  The joint automated repository for various integrated simulations (jarvis) for data-driven materials design.  npj computational materials, 6(1):173, 2020.
  * [53] Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky.  Learning local equivariant representations for large-scale atomistic dynamics.  arXiv preprint arXiv:2204.05249, 2022.
  * [54] Tuan Le, Frank Noé, and Djork-Arné Clevert.  Equivariant graph attention networks for molecular property prediction, 2022.
  * [55] S.D. Bond and B.J. Leimkuhler.  Molecular dynamics and the accuracy of numerically computed averages.  Acta Numerica, 16(-1):1–65, 2007.
  * [56] Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, E Weinan, and Linfeng Zhang.  Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning.  In SC20: International conference for high performance computing, networking, storage and analysis, pages 1–14. IEEE, 2020.
  * [57] Meng-Hao Guo, Tian-Xing Xu, Jiang-Jiang Liu, Zheng-Ning Liu, Peng-Tao Jiang, Tai-Jiang Mu, Song-Hai Zhang, Ralph R Martin, Ming-Ming Cheng, and Shi-Min Hu.  Attention mechanisms in computer vision: A survey.  Computational Visual Media, pages 1–38, 2022.
  * [58] Andrea Galassi, Marco Lippi, and Paolo Torroni.  Attention in natural language processing.  IEEE Transactions on Neural Networks and Learning Systems, 32(10):4291–4308, 2020.
  * [59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [60] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [61] Jianxing Huang, Linfeng Zhang, Han Wang, Jinbao Zhao, Jun Cheng, and Weinan E.  Deep potential generation scheme and simulation protocol for the li10gep2s12-type superionic conductors.  The Journal of Chemical Physics, 154(9):094703, 2021.
  * [62] Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller.  Machine learning of accurate energy-conserving molecular force fields.  Science advances, 3(5):e1603015, 2017.
  * [63] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld.  Quantum chemistry structures and properties of 134 kilo molecules.  Scientific data, 1(1):1–7, 2014.
  * [64] Yifei Mo, Shyue Ping Ong, and Gerbrand Ceder.  First principles study of the li10gep2s12 lithium super ionic conductor material.  Chemistry of Materials, 24(1):15–17, 2012.
  * [65] Aris Marcolongo and Nicola Marzari.  Ionic correlations and failure of nernst-einstein relation in solid-state electrolytes.  Physical Review Materials, 1(2):025402, 2017.
  * [66] Alexander Kuhn, Jürgen Köhler, and Bettina V Lotsch.  Single-crystal x-ray structure analysis of the superionic conductor li 10 gep 2 s 12.  Physical Chemistry Chemical Physics, 15(28):11620–11622, 2013.

##  Appendix A  Ablation study

Table A.1:  Ablation study of the DPA-1 model architecture. The models are validated by the energy (meV/atom) and atomic forces (meV/Å) RMSEs calculated with different settings of the training/validation sets. Bold numbers correspond to lowest values.  | HEA validation RMSE
---|---
Training | interior | exterior | all
Validation | exterior | interior | all
|  | Energy
---
(meV/atom)
| Force
---
(meV/Å)
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
DeepPotSE | 141.32 | 257 | 40.74 | 204 | 15.75 | 117
| DPA-1 w/o attention
---
(tebd)
88.07 | 252 | 23.66 | 157 | 7.69 | 88
| DPA-1 w/o gated
---
(tebd+attention)
34.79 | 191 | 9.59 | 120 | 6.35 | 75
| DPA-1 gaussian
---
(tebd+attention+gated)
35.55 | 161 | 7.01 | 121 | 4.97 | 77
| DPA-1 full
---
(tebd+attention+gated)
31.27 | 158 | 6.88 | 117 | 4.96 | 71

We conducted an ablation study on the model structure of DPA-1 for HEA
systems, as illustrated in Table A.1. Our focus was on three primary
components that make the DPA-1 different from the DeepPot-SE model: type
embedding defined by Eq. 4 (denoted as tebd), attention updates by Eq. 9
(denoted as attention), and the gate of
$\mathrm{\hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T}}$ in Eq. 8 (denoted
as gate). The full DPA-1 structure is thus “tebd+attention+gate”. Our findings
indicate that tebd and attention substantially enhance performance,
potentially due to the effective encoding of elemental space and the
communication of information within the neighborhood. The gate component in
the attention map encodes the angular information of neighboring atoms. It
provides a relatively smaller but not negligible contribution to the model
accuracy. Moreover, we examined the impact of different encodings of the
radial information by altering the neural network defined embedding (i.e. Eq.
(5)) to a Gaussian kernels defined embedding (denoted by DPA-1 gaussian). Note
that the distance between atoms are not inversed and scaled as Eq. 3, but
affine transformed with element-type-dependent parameters [37]. A slight but
observable reduction in the test accuracy is observed. This change renders the
model non-smooth under the periodic boundary condition, i.e. a sudden energy
jump may occur as the positions of atoms vary infinitesimally. This leads to a
non-conserved energy in the Hamiltonian dynamics simulations.

##  Appendix B  Equilibrium properties of tenary AlMgCu alloys

We performed geometry relaxation calculations on all AlMgCu ternary alloys
available from the Materials Project [49]. We compared single-binary-trained
and fully-trained DPA-1 and DeepPotSE models on the accuracy of formation
energy and equilibrium volume. Mean absolute errors (MAEs) in Fig A.1 (a)
shows that DPA-1 model trained only on single and binary samples, achieves
remarkable enhancement in accuracy on ternary alloys compared with the single-
binary-trained DeepPot-SE, confirming DPA-1’s superior generalization
capabilities. The enhancement is around one order of magnitude for both
properties, from about 400 meV/atom to less than 40 meV/atom in formation
energy and from about 1 Å3/atom to about 0.1 Å3/atom in equilibrium volume.
Structure-resolved errors in Fig A.1 (b)(c) clearly indicate that, the
performance of the DPA-1 model on all six ternary alloys, is comparable with
that of the fully trained DPA-1 or DeepPot-SE.

Figure A.1:  (a) Mean absolute errors (MAEs) of formation energy (FormE) and
equilibrium volume (EquiV) for all AlMgCu ternary alloys from the Materials
Project, as predicted by single-binary-trained and fully-trained DPA-1 and
DeepPot-SE models after relaxation. (b) (c) Structure-resolved absolute errors
(AEs).

##  Appendix C Test on the elastic moduli of AlMgCu alloys

Test results of elastic moduli of binary and tenary AlMgCu alloys are shown in
Fig. A.2, with the overall RMSE of both moduli are respectively 7.41, 5.41 and
7.56 GPa for DeepPot-SE, DPA-1 and 10%-DPA-1. Finite difference method is used
to calculate the elastic moduli, where the maximum norm and shear deformation
with respect to the equilibrium conformation are both 2%.

Figure A.2: Test of trained DeepPot-SE and DPA-1 models on shear (a) and bulk
(b) moduli of AlMgCu systems. Both DeepPot-SE and DPA-1 models are trained on
all subsets of AlMgCu systems, while 10%-DPA-1 model used less ternary samples
(10%).

##  Appendix D  Diffusion coefficients on LiGePS systems

Figure A.3:  Diffusion coefficients of Li versus temperature in the
Li10GeP2S12 system obtained from DPA-1 and DeepPot-SE [61] models trained on
SSE systems, previous AIMD simulations [64, 65] and experimental solid-state
NMR results [66].

We carried out molecular dynamics simulations on LiGePS systems to assess the
diffusion coefficients of Li in relation to temperature, akin to the approach
in [61]. As illustrated in Fig. A.3, the DPA-1 models trained on SSE systems
exhibited a high degree of consistency when compared to AIMD simulations in
[64, 65], experimental studies in [66] and DeepPot-SE simulations in [61].
This consistency lends credibility to their reliability.

##  Appendix E  Embedded element type interpolation for Nb and Mo in the OC2M
dataset

Figure A.4:  The RMSE of energy and force on the HEA system using embedded
element type interpolation for (a)(b) single-element replacements and (c)(d)
two-element replacements. Nbwo, Mowo and NbMowo denote models trained on
purposefully altered OC2M datasets that excluded frames containing Nb, Mo, and
either Nb or Mo, respectively.

To demonstrate the interpretability of DPA-1, we trained models with certain
elements omitted from the OC2M dataset and conducted element type
interpolation experiments in a manner similar to that described in Sec. III.4.

In the first two experiments, DPA-1 models named Nbwo, Mowo are trained on the
OC2M datasets excluding the data frames containing Nb, Mo elements,
respectively. We used the learned type embeddings of elements from the same
periodic table family to interpolate omitted elements as replacements, given
by

$\displaystyle
T_{Nb}=\lambda\left(V\right)*T_{V}+\left(1-\lambda\left(V\right)\right)*T_{Ta}$
(13) $\displaystyle
T_{Mo}=\lambda\left(Cr\right)*T_{Cr}+\left(1-\lambda\left(Cr\right)\right)*T_{W},$
(14)

respectively. As shown in Fig. A.4 (a)–(b), we evaluated the RMSE of energy
and force on bulk TaNbWMoVAl HEA systems described in Sec. III.1. The Nbwo
model achieves an optimal accuracy of 14.2 meV/atom (energy) and 236 meV/Å
(force) at $\lambda\left(V\right)=0.5$, while the Mowo model gives the lowest
energy error of 12.3 meV/atom and force error of 218 meV/Å at
$\lambda\left(Cr\right)=0.5$. The optimal accuracy of Nbwo and Mowo models are
almost the same as the baseline errors of 12.0 meV/atom (energy) and 232 meV/Å
(force), given by a full OC2M trained DPA-1 model tested on the HEA dataset.
Surprisingly, the Mowo model presents a higher force accuracy than the
baseline model.

We further remove all data frames containing either Nb or Mo elements, and
train a DPA-1 model named NbMowo. The type representation of Nb and Mo are
obtained by the interpolation of type embeddings of V, Ta, Cr and W, i.e. Eqs.
(13)–(14). As shown in Fig. A.4 (c)–(d), the optimal interpolation weights for
Nb and Mo are $\lambda\left(V\right)=0.4$ and $\lambda\left(Cr\right)=0.6$,
respectively. The corresponding energy error is 13.4 meV/atom lying between
the Nbwo and Mowo models. The force error 253 meV/Å is slightly higher than
the Nbwo and Mowo models, but is only 10% higher than the baseline. The
interpolation experiments of Nb and Mo further validate the interpretability
of the DPA-1 element type embedding.

##  Appendix F Results on simple datasets

Table A.2: Prediction RMSE of energies (meV/atom) and atomic forces (meV/Å) given by DPA-1 (OC2M-pretrained), DPA-1 (from scratch), EANN, and DeepPot-SE. DPA-1 (OC2M-pretrained) means that the results were obtained by funetuning a DPA-1 model pretrained using the OC2M dataset. For training DPA-1, EANN, and DeepPot-SE, 10%, 15$\sim$20%, and 90% randomly sampled data points are used, respectively. Bold numbers correspond to lowest values. Systems | Sub-systems |  | DPA-1
---
(OC2M-pretrained)
DPA-1 | EANN | DeePot-SE
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
| Energy
---
(meV/atom)
| Force
---
(meV/Å)
Cu |  | FCC
---
solid
0.12 | 86 | 0.14 | 88 | 0.16 | 89 | 0.18 | 90
Ge |  | diamond
---
solid
0.09 | 21 | 0.10 | 24 | 0.07 | 31 | 0.35 | 38
Si |  | diamond
---
solid
0.09 | 23 | 0.11 | 25 | 0.08 | 28 | 0.24 | 36
Al2O3 |  | Trigonal
---
solid
0.08 | 40 | 0.08 | 41 | 0.11 | 51 | 0.23 | 49
C5H5N | Pyridine-I | 0.17 | 21 | 0.21 | 23 | 0.17 | 41 | 0.38 | 25
Pyridine-II | 0.24 | 29 | 0.26 | 32 | 0.35 | 49 | 0.65 | 39
TiO2 | Rutile | 0.56 | 109 | 0.62 | 112 | 0.51 | 133 | 0.96 | 137
Anatase | 0.87 | 144 | 0.91 | 151 | 1.03 | 183 | 1.78 | 181
Brookite | 0.50 | 88 | 0.52 | 91 | 0.55 | 96 | 0.59 | 94
MoS2 +Pt | MoS2 slab | 0.14 | 11 | 0.17 | 12 | 0.26 | 19 | 5.26 | 23
bulk Pt | 0.25 | 41 | 0.26 | 43 | 0.38 | 64 | 2.00 | 84
Pt surface | 0.59 | 45 | 0.68 | 45 | 3.8 | 86 | 6.77 | 105
Pt cluster | 1.21 | 41 | 1.33 | 41 | 9.7 | 152 | 30.6 | 201
Pt on MoS2 | 4.43 | 86 | 5.02 | 88 | 1.33 | 91 | 2.62 | 94
CoCrFe MnNi | rand.occ.I | 1.31 | 367 | 1.43 | 387 | 2.3 | 410 | 1.68 | 394
rand.occ.II | 1.74 | 381 | 1.93 | 389 | 3.3 | 415 | 5.29 | 410

As shown in Table A.2, we trained DPA-1 from scratch on simple bulk systems
[6] and compared with the embedded atom neural network potential (EANN) [10]
and DeepPot-SE.

This small dataset contains two types of systems. The first type includes
general systems, such as relatively easy Cu, Ge, Si, Al_2O_3 with one single
solid phase, and more challenging systems like C_5H_5N (pyridine) and TiO_2
with two and three phases, respectively. The second type of systems contains a
grand-canonical-like system of supported Pt clusters on a MoS_2 slab and a
CoCrFeMnNi high entropy alloy (HEA) system.

For training DPA-1 and EANN, only 10% and 15$\sim$20% randomly selected data
points are used respectively, while DeepPot-SE used 90%. The results show
that, even with less training samples and trained from scratch, DPA-1 still
mostly outperforms the other methodologies, especially in terms of force
prediction accuracy. Interestingly, all the finetuned model based on
OC2M-pretrained outperforms other models.

##  Appendix G Hyperparameters in sample efficiency tests

To generate the learning curves of sample efficiency in Sec. III.3, we started
from 50 randomly selected samples and set the training steps to be 100k during
each iteration. The probabilities to visit the historical and newly added
samples are set to 0.9 and 0.1, respectively.