Environment-adaptive machine learning potentials

We consider a multi-element system of $N_{a}$ atoms with $N_{e}$ unique elements. We denote by $\bm{r}_{i}$ and $Z_{i}$ position vector and type of an atom $i$ in the system, respectively. Thus we have $Z_{i}\in\{1,\ldots,N_{e}\}$, $\bm{R}=(\bm{r}_{1},\bm{r}_{2},\ldots,\bm{r}_{N_{a}})\in\mathbb{R}^{3N_{a}}$, and $\bm{Z}=(Z_{1},Z_{2},\ldots,Z_{N_{a}})\in\mathbb{N}^{N_{a}}$. The potential energy surface (PES) of the system of $N_{a}$ atoms can be expressed as a many-body expansion of the form

$$\begin{split}E_{\rm T}(\bm{R},\bm{Z})=&\sum_{i}V^{(1)}(\bm{r}_{i},Z_{i})+\sum_{i,j}V^{(2)}(\bm{r}_{i},\bm{r}_{j},Z_{i},Z_{j})\\ &+\sum_{i,j,k}V^{(3)}(\bm{r}_{i},\bm{r}_{j},\bm{r}_{k},Z_{i},Z_{j},Z_{k})\\ &+\sum_{i,j,k,l}V^{(4)}(\bm{r}_{i},\bm{r}_{j},\bm{r}_{k},\bm{r}_{l},Z_{i},Z_{j},Z_{k},Z_{l})+\ldots\end{split}$$

(1)

The superscript on each potential denotes its body order. Each potential must also depend on a set of parameters used to parametrize it for a specific application. To simplify the notation, we have chosen not to explicitly denote these parameters in the potentials. A separation of the PES into atomic contributions yields

$$E_{\rm T}(\bm{R},\bm{Z})=\sum_{i=1}^{N_{a}}E_{i}(\bm{R},\bm{Z})$$

(2)

where $E_{i}$ is obtained from (1) by removing the sum over index $i$. To make the PES invariant with respect to translation and rotation, the potentials should depend only on internal coordinates as follows

$$\begin{split}E_{i}=&V^{(1)}(Z_{i})+\sum_{j}V^{(2)}(r_{ij},Z_{i},Z_{j})\ +\\ &\sum_{j,k}V^{(3)}(r_{ij},r_{ik},w_{ijk},Z_{i},Z_{j},Z_{k})\ +\\ &\sum_{j,k,l}V^{(4)}(r_{ij},r_{ik},r_{il},w_{ijk},w_{ijl},w_{ikl},Z_{i},Z_{j},Z_{k},Z_{l})+\ldots\end{split}$$

(3)

where $\bm{r}_{ij}=\bm{r}_{j}-\bm{r}_{i}$, $r_{ij}=|\bm{r}_{ij}|$, $w_{ijk}=\cos\theta_{ijk}=\hat{\bm{r}}_{ij}\cdot\hat{\bm{r}}_{ik}$, $\hat{\bm{r}}=\bm{r}/|\bm{r}|$. The internal coordinates include both distances $r_{ij},r_{ik},r_{il}$ and angles $w_{ijk},w_{ijl},w_{ikl}$. The number of internal coordinates for $V^{(q)}$ is equal to $(q-1)q/2$. Typically, the one-body terms $V^{(1)}(Z_{i})$ are set to the isolated energies of atom $i$.

We briefly describe the construction of two-body PODs and refer to [17, 26] for further details. We assume that the direct interaction between two atoms vanishes smoothly when their distance is greater than the cutoff distance $r_{\rm cut}$. Furthermore, we assume that two atoms can not get closer than the inner cutoff distance $r_{\rm in}$. Letting $r\in(r_{\rm in},r_{\rm cut})$, we introduce the following parametrized radial functions

$$\phi(r,r_{\rm in},r_{\rm cut},\alpha,\beta)=\frac{\sin(\alpha\pi x)}{\alpha(r-r_{\rm in})},\qquad\varphi(r,\gamma)=\frac{1}{r^{\gamma}},$$

(4)

where the scaled distance function $x$ is given by

$$x(r,r_{\rm in},r_{\rm cut},\beta)=\frac{e^{-\beta(r-r_{\rm in})/(r_{\rm cut}-r_{\rm in})}-1}{e^{-\beta}-1}.$$

(5)

The function $\phi$ in (4) is related to the zeroth spherical Bessel function, while the function $\varphi$ is inspired by the n-m Lennard-Jones potential. Although the parameter $\gamma$ can be real number, we choose a set of consecutive positive integers $\{1,2,\ldots,P_{\gamma}\}$ to compute instances of the parametrized function $\varphi$ by making use of the relation $\varphi(r,\gamma+1)=\varphi(r,\gamma)/r$. Similarly, we choose a set of consecutive integers $\{1,2,\ldots,P_{\alpha}\}$ for $\alpha$ to generate instances of the parametrized function $\phi$ by making use of the formula $\sin((\alpha+1)\pi x)=\sin(\pi x)U_{\alpha}(\cos(\pi x))$, where $U_{\alpha}$ are Chebyshev polynomials of the second kind. We take $P_{\beta}$ values for the parameter $\beta$ such that $\beta_{k}=(k-1)\beta_{\max}/(P_{\beta}-1)$ for $k=1,2,\ldots,P_{\beta}$, where $\beta_{\max}=4.0$.

We introduce the following function as a convex combination of the two functions in (4)

$$\psi(r,\bm{\mu})=\kappa\phi(r,r_{\rm in},r_{\rm cut},\alpha,\beta)+(1-\kappa)\varphi(r,\gamma),$$

(6)

where $\mu_{1}=r_{\rm in},\mu_{2}=r_{\rm cut},\mu_{3}=\alpha,\mu_{4}=\beta,\mu_{5}=\gamma$, and $\mu_{6}=\kappa$. The two-body parametrized potential is defined as follows

$$V^{(2)}(r_{ij},\bm{\mu})=f_{\rm c}(r_{ij},\bm{\mu})\psi(r_{ij},\bm{\mu})$$

(7)

where the cut-off function $f_{\rm c}(r_{ij},\bm{\mu})$ is

$$f_{\rm c}(r,\bm{\mu})=\exp\left(1-\frac{1}{\sqrt{\left(1-\frac{(r-r_{\rm in})^{3}}{(r_{\rm cut}-r_{\rm in})^{3}}\right)^{2}+\epsilon}}\right)$$

(8)

with $\epsilon=10^{-6}$. This cut-off function ensures the smooth vanishing of the two-body potential and its derivative for $r_{ij}\geq r_{\rm cut}$.

Given $S=P_{\gamma}+P_{\alpha}P_{\beta}$ parameter tuples $\bm{\mu}_{s},1\leq s\leq S$, we introduce the following set of snapshots

$$\Phi_{s}(r_{ij})=V^{(2)}(r_{ij},\bm{\mu}_{s}),\quad s=1,\ldots,S.$$

(9)

We next employ the proper orthogonal decomposition [17] to generate an orthogonal basis set which is known to be optimal for representation of the snapshot family $\{\Phi_{s}\}_{s=1}^{S}$. In particular, the orthogonal radial basis functions are computed as follows

$$R_{n}(r_{ij})=\sum_{s=1}^{S}Q_{sn}\,\Phi_{s}(r_{ij}),\qquad n=1,\ldots,N_{r},$$

(10)

where the number of radial basis functions $N_{r}$ is typically in the range between 5 and 10. Note that $Q_{sn},1\leq s\leq S,1\leq n\leq N_{r},$ are a matrix whose columns are eigenvectors of the following eigenvalue problem

$$\bm{C}\bm{a}=\lambda\bm{a},$$

(11)

where the covariance matrix $\bm{C}$ is given by

$$C_{sp}=\int_{r_{\rm in}}^{r_{\rm cut}}\Phi_{s}(r)\Phi_{p}(r)dr,\qquad 1\leq s,p\leq S.$$

(12)

The covariance matrix is computed by using the trapezoidal rule on a grid of 2000 subintervals on the interval $[r_{\rm in},r_{\rm cut}]$. The eigenvector matrix $Q_{sn}$ is pre-computed and stored.

Finally, the two-body proper orthogonal descriptors at each atom $i$ are computed by summing the orthogonal basis functions over the neighbors of atom $i$ and numerating on the atom types as follows

$$\mathcal{D}^{(2)}_{ipqn}=\left\{\begin{array}[]{ll}\displaystyle\sum_{\{j=1|Z_{j}=q\}}^{N_{i}}R_{n}(r_{ij}),&\quad\mbox{if }Z_{i}=p\\ 0,&\quad\mbox{if }Z_{i}\neq p\end{array}\right.$$

(13)

for $1\leq i\leq N_{a},1\leq n\leq N_{r},1\leq q,p\leq N_{e}$. The number of two-body descriptors per atom is thus $N_{r}N_{e}^{2}$.

For the purpose of complexity analysis, we assume that each atom has the same number of neighbors $N_{i}$. The total number of neighbors is thus $N_{a}N_{i}$ for all atoms. The cost of evaluating the radial basis functions in (10) is $O(N_{a}N_{i}N_{r}S)$, while the cost of evaluating the two-body descriptors in (13) is $O(N_{a}N_{i}N_{r})$. The total cost is thus independent of the number of elements $N_{e}$.

For any given integer $\ell\in[0,P_{a}]$, where $P_{a}$ is the highest angular degree, we introduce a basis set of angular monomials

$$A_{\ell m}(\hat{\bm{r}}_{ij})=(\hat{x}_{ij})^{l_{x}}\ (\hat{y}_{ij})^{l_{y}}\ (\hat{z}_{ij})^{l_{z}},$$

(14)

where the exponents $l_{x},l_{y},l_{z}$ are nonnegative integers such that ${l_{x}}+l_{y}+l_{z}=\ell$. Note that the index $m$ satisfies $0\leq m\leq(\ell+1)(\ell+2)/2-1$ for any given $\ell$, and that the total number of angular monomials is $(P_{a}+1)(P_{a}+2)(P_{a}+3)/6$. Table 1 shows the basis set of angular monomials for $P_{a}=4$. By applying the trinomial expansion to the power of the angle component $w_{ijk}$, we obtain

$$\begin{split}(w_{ijk})^{\ell}&=\left(\hat{x}_{ij}\hat{x}_{ik}+\hat{y}_{ij}\hat{y}_{ik}+\hat{z}_{ij}\hat{z}_{ik}\right)^{\ell}\\ &=\sum_{m=0}^{L(\ell)}C_{\ell m}A_{\ell m}(\hat{\bm{r}}_{ij})A_{\ell m}(\hat{\bm{r}}_{ik})\end{split}$$

(15)

where $L(\ell)=(\ell+1)(\ell+2)/2-1$ and $C_{\ell m}=\frac{\ell!}{l_{x}!l_{y}!l_{z}!}$ correspond to the multinomial coefficients of the trinomial expansion. Table 2 shows the multinomial coefficients for $\ell=0,1,2,3,4$.

Next, we define the atom basis functions at atom $i$ as the sum over all neighbors of atom $i$ of the products of radial basis functions and angular monomials

$$B_{iqn\ell m}=\sum_{j=1|Z_{j}=q}^{N_{i}}R_{n}(r_{ij})A_{\ell m}(\hat{\bm{r}}_{ij}),\quad 1\leq q\leq N_{e}.$$

(16)

The cost of evaluating the atom basis functions is $O(N_{a}N_{i}N_{r}(P_{a}+1)(P_{a}+2)(P_{a}+3)/6)$, which is independent of the number of elements. These atom basis functions are used to define the atom density descriptors as follows

$$\mathcal{D}^{(3)}_{ipqq^{\prime}n\ell}=\left\{\begin{array}[]{ll}\displaystyle\sum_{m=0}^{L(\ell)}C_{\ell m}B_{iqn\ell m}B_{iq^{\prime}n\ell m},&\quad\mbox{if }Z_{i}=p\\ 0,&\quad\mbox{if }Z_{i}\neq p\end{array}\right.$$

(17)

for $1\leq i\leq N_{a},1\leq n\leq N_{r},0\leq\ell\leq P_{a},1\leq q,p\leq N_{e},1\leq q^{\prime}\leq q$. The number of three-body descriptors per atom is thus $N_{r}(P_{a}+1)N_{e}^{2}(N_{e}+1)/2$. The cost of evaluating (17) is $O(N_{a}N_{r}N_{e}(N_{e}+1)(P_{a}+1)(P_{a}+2)(P_{a}+3)/12)$, which is usually less than that of evaluating the atom basis functions. Therefore, the total cost of computing the atom density descriptors is $O(N_{a}N_{i}N_{r}(P_{a}+1)(P_{a}+2)(P_{a}+3)/6)$.

Substituting (16) into (17) yields the internal coordinate form of the three-body descriptors

$$\mathcal{D}^{(3)}_{ipqq^{\prime}n\ell}=\displaystyle\sum_{j|Z_{j}=q}^{N_{i}}\sum_{k|Z_{k}=q^{\prime}}^{N_{i}}R_{n}(r_{ij})R_{n}(r_{ik})\left(w_{ijk}\right)^{\ell}$$

(18)

for $Z_{i}=p$. The cost of evaluating the internal coordinate descriptors is $O(N_{a}N^{2}_{i}N_{r}P_{a})$. Although the atom density descriptors (17) and the internal coordinate descriptors (18) are mathematically equivalent, their computational complexity are not the same. The complexity of the atom density descriptors is linear in the number of neighbors and cubic in the angular degree, whereas that of the internal coordinate descriptors is quadratic in the number of neighbors and linear in the angular degree. Therefore, if the number of neighbors is large and the angular degree is small, it is faster to evaluate the the atom density descriptors. On the other hand, if the number of neighbors is small and the angular degree is high, it is more efficient to compute the internal coordinate descriptors.

We begin by introducing the following four-body angular functions

$$f_{s}(w_{ijk},w_{ijl},w_{ikl})=\left(w_{ijk}\right)^{a}\left(w_{ijl}\right)^{b}\left(w_{ikl}\right)^{c},$$

(19)

where $a,b,c$ are integers such that $a+b+c=\ell$ and $a\geq b\geq c\geq 0$. The four-body angular functions $f_{s}$ are listed in Table 3. The four-body internal coordinate descriptors at each atom $i$ are defined as

$$\mathcal{D}^{(4)}_{ipqq^{\prime}q^{\prime\prime}ns}=\displaystyle\sum_{\{j|Z_{j}=q\}}^{N_{i}}\sum_{\{k|Z_{k}=q^{\prime}\}}^{N_{i}}\sum_{\{l|Z_{l}=q^{\prime\prime}\}}^{N_{i}}U_{ns}$$

(20)

for $Z_{i}=p$, where $U_{ns}$ are given by

$$U_{ns}=R_{n}(r_{ij})R_{n}(r_{ik})R_{n}(r_{il})f_{s}(w_{ijk},w_{ijl},w_{ikl})$$

(21)

for $1\leq i\leq N_{a},1\leq n\leq N_{r},1\leq s\leq K_{a},1\leq p,q\leq N_{e},1\leq q^{\prime}\leq q,1\leq q^{\prime\prime}\leq q^{\prime}$. Here $K_{a}$ is the number of four-body angular basis functions, which depends on $P_{a}$. The number of four-body descriptors per atom is thus $N_{r}K_{a}N_{e}^{2}(N_{e}+1)(N_{e}+2)/6$. The cost of evaluating the four-body internal coordinate descriptors is $O(N_{a}N^{3}_{i}N_{r}K_{a})$, which is independent of the number of elements.

We note from the trinomial expansion that

$$\begin{split}(\xi_{1}&+\xi_{2}+\xi_{3})^{a}(\eta_{1}+\eta_{2}+\eta_{3})^{b}(\zeta_{1}+\zeta_{2}+\zeta_{3})^{c}=\\ &\sum_{\alpha=0}^{L(a)}\sum_{\beta=0}^{L(b)}\sum_{\gamma=0}^{L(c)}C_{a\alpha}C_{b\beta}C_{c\gamma}A_{a\alpha}(\bm{\xi})A_{b\beta}(\bm{\eta})A_{c\gamma}(\bm{\zeta})\end{split}$$

where $A_{a\alpha}$ are the angular monomials defined in (14). By considering $\xi_{1}=\hat{x}_{ij}\hat{x}_{ik},\xi_{2}=\hat{y}_{ij}\hat{y}_{ik},\xi_{3}=\hat{z}_{ij}\hat{z}_{ik}$, $\eta_{1}=\hat{x}_{ij}\hat{x}_{il},\eta_{2}=\hat{y}_{ij}\hat{y}_{il},\eta_{3}=\hat{z}_{ij}\hat{z}_{il}$, $\zeta_{1}=\hat{x}_{ik}\hat{x}_{il},\zeta_{2}=\hat{y}_{ik}\hat{y}_{il},\zeta_{3}=\hat{z}_{ik}\hat{z}_{il}$, we obtain

$$A_{a\alpha}(\bm{\xi})A_{b\beta}(\bm{\eta})A_{c\gamma}(\bm{\zeta})=A_{a^{\prime}\alpha^{\prime}}(\hat{\bm{r}}_{ij})A_{b^{\prime}\beta^{\prime}}(\hat{\bm{r}}_{ik})A_{c^{\prime}\gamma^{\prime}}(\hat{\bm{r}}_{il})$$

where $a^{\prime}=a+b$, $b^{\prime}=a+c$, $c^{\prime}=b+c$, and the index $\alpha^{\prime}$ depends on $\alpha$ and $\beta$, $\beta^{\prime}$ on $\alpha$ and $\gamma$, $\gamma^{\prime}$ on $\beta$ and $\gamma$. It thus follows that the four-body angular functions can be expressed as

$$f_{s}=\sum_{\alpha=0}^{L(a)}\sum_{\beta=0}^{L(b)}\sum_{\gamma=0}^{L(c)}C_{abc}^{\alpha\beta\gamma}A_{a^{\prime}\alpha^{\prime}}(\hat{\bm{r}}_{ij})A_{b^{\prime}\beta^{\prime}}(\hat{\bm{r}}_{ik})A_{c^{\prime}\gamma^{\prime}}(\hat{\bm{r}}_{il}).$$

where $C_{abc}^{\alpha\beta\gamma}=C_{a\alpha}C_{b\beta}C_{c\gamma}$. Hence, the internal coordinate form (20) is equivalent to the atom density form

$$\begin{split}\mathcal{D}^{(4)}_{ipqq^{\prime}q^{\prime\prime}ns}=&\sum_{\alpha=0}^{L(a)}\sum_{\beta=0}^{L(b)}\sum_{\gamma=0}^{L(c)}C_{abc}^{\alpha\beta\gamma}B_{iqna^{\prime}\alpha^{\prime}}B_{iq^{\prime}nb^{\prime}\beta^{\prime}}B_{iq^{\prime\prime}nc^{\prime}\gamma^{\prime}}.\end{split}$$

(22)

The four-body atom density descriptors are expressed in terms of the sums of the products of the atom basis functions. In general, they are more efficient to evaluate than their internal coordinate counterparts because they scale linearly with the number of neighbors. The complexity analysis of the four-body atom density descriptors is detailed in [26].

It is possible to exploit the symmetry and hierarchy of the four-body descriptors to reduce the computational cost. In particular, the four-body descriptors associated with $f_{1}=1$ are the products of three two-body descriptors. For $b=c=0$ (e.g., $f_{2},f_{3},f_{5},f_{8}$ in Table 3) the four-body descriptors are the products of two-body descriptors and three-body descriptors. Those four-body descriptors can be computed very fast without using the atom density form (22). The remaining four-body descriptors are calculated by using (22). They share many common terms, which can be exploited to further reduce the cost. For instance, since $f_{6}=w_{ijk}f_{4}$ and $f_{9}=w^{2}_{ijk}f_{4}$, the descriptors associated with $f_{6}$ and $f_{9}$ have many common terms with those associated with $f_{4}$.

Linear regression is the most simple and efficient method for building MLIP models [21, 27, 28, 22]. Let ${\mathcal{D}}_{im},1\leq m\leq M,$ be a set of $M$ local descriptors at atom $i$. The atomic energy at an atom $i$ is expressed as a linear combination of the local descriptors

$$E_{i}(\bm{R},\bm{Z})=\sum_{m=1}^{M}c_{m}\mathcal{D}_{im}(\bm{R},\bm{Z})$$

(23)

where $c_{m}$ are the coefficients to be determined by fitting against QM database. The PES is given by

$$E_{\rm T}(\bm{R},\bm{Z})=\sum_{i=1}^{N_{a}}E_{i}(\bm{R},\bm{Z})=\sum_{m=1}^{M}c_{m}\mathcal{G}_{m}(\bm{R},\bm{Z})$$

(24)

where $\mathcal{G}_{m}=\sum_{i=1}^{N_{a}}\mathcal{D}_{im}$ are the global descriptors. The coefficients $c_{m}$ are sought as solution of a least squares problem

$$\min_{\bm{c}}\alpha\|\bm{G}\bm{c}-\bm{e}\|^{2}+\beta\|\bm{H}\bm{c}-\bm{f}\|^{2}+\gamma\|\bm{c}\|^{2}$$

(25)

where the matrix $\bm{G}$ is formed from the global descriptors, while $\bm{H}$ is formed from the derivatives of the global descriptors for all configurations in the training database. The vector $\bm{e}$ is comprised of DFT energies, while $\bm{f}$ is comprised of DFT forces. Note that $\alpha$ is the energy weight parameter, $\beta$ is the force weight parameter, and $\gamma$ is a regularization parameter. They are hyperparameters of the linear regression model.

In order to accurately model atomic forces in MD simulations, the training dataset must be diverse and rich enough to cover structural forms (e.g., crystalline, amorphous, defects, interfaces), multiple phases (solid, liquid, gas, plasma), and a wide range of temperature, pressure, and chemical environment. Training a linear model on the entire training set may not produce an accurate and efficient potential it the dataset contains very diverse environments. We partition the training dataset into $K$ separate subsets, whose DFT energies and forces are denoted by $(\bm{e}^{k},\bm{f}^{k}),1\leq k\leq K$. On each subset, we introduce an associated PES

$$E^{k}_{\rm T}(\bm{R},\bm{Z})=\sum_{i=1}^{N_{a}}\sum_{m=1}^{M}c_{m}^{k}\mathcal{D}_{im}(\bm{R},\bm{Z}),$$

(26)

where the coefficient vectors $\bm{c}^{k}$ are sought as solutions of the least squares problems

$$\min_{\bm{c}^{k}}\alpha^{k}\|\bm{G}^{k}\bm{c}^{k}-\bm{e}^{k}\|^{2}+\beta^{k}\|\bm{H}^{k}\bm{c}^{k}-\bm{f}^{k}\|^{2}+\gamma^{k}\|\bm{c}^{k}\|^{2}.$$

(27)

Here the superscript $k$ is used to indicate the quantities associated with the $k$th subset. This training strategy yields an ensemble of $K$ separate potentials. Each potential may accurately predict configurations in the subset on which it is trained. However, it may not be accurate for predicting configurations in the other subsets.

A hypothetical issue here is to guarantee the continuity of PES when predicting forces with an ensemble of potentials. It is not obvious how to combine these separate potentials, as they are trained on different datasets. A simple strategy is to select the best potential among these potentials to predict the physical properties of a given configuration at hand, if the criterion of selecting the best potential can be defined. While this strategy may work for property prediction, it does not work for MD simulations. This is because using different potentials in an MD simulation will result in discontinuity in the PES and thus forces. The remainder of this section describes a method that allows us to combine these separate potentials to construct a global, differentiable and continuous PES.

We describe a clustering method to partition the dataset into subsets of similar attributes. Local atomic environment of an atom $i$ comprises the positions and chemical species of the atom and its neighbors within a cutoff radius. These atom positions and chemical species can be mapped onto a vector of $M$ invariant descriptors, $\mathcal{D}_{im},1\leq m\leq M$, by using either the internal coordinate approach or the atom density approach. For a dataset of $N$ atoms, we obtain a descriptor matrix $\bm{\mathcal{D}}$ of size $N$ by $M$. Each row of the descriptor matrix encapsulates the local atomic environment of the corresponding atom. The similarity between two local atomic environments can be measured by the dot product of the two corresponding descriptor vectors. Therefore, partitioning the dataset of $N$ atoms into $K$ clusters can be done by dividing the rows of the descriptor matrix into $K$ similarity subsets. One can use a clustering method such as $k$-means clustering algorithm to partition $N$ vectors in $M$ dimensions into $K$ separate clusters. However, clustering in very high dimensions can be very expensive, since $M$ is typically large.

To reduce computational cost, we consider a dimensionality reduction method to compress the descriptor matrix into a lower-dimensional matrix of size $N$ by $J$, where $J$ is considerably less than $M$. In this paper, principal component analysis is used to obtain the low-dimensional descriptor matrix as follows

$$\bm{\mathcal{B}}=\bm{\mathcal{D}}\ \bm{W}\ .$$

(28)

Here $\bm{W}\in\mathbb{R}^{M\times J}$ consists of the first $J$ eigenvectors of the eigenvalue decomposition $\bm{\mathcal{D}}^{T}\bm{\mathcal{D}}=\bm{W}\bm{\Lambda}\bm{W}^{T}$, and the eigenvalues are ordered from the largest to the smallest.

For multi-element systems, the clustering method is applied to the local descriptor matrix for each element as follows. The matrix $\bm{\mathcal{D}}$ is split into $\bm{\mathcal{D}}^{e},1\leq e\leq N_{e}$, where $\bm{\mathcal{D}}^{e}$ is formed by gathering the rows of $\bm{\mathcal{D}}$ for all atoms of the element $e$. For each element $e$, we compute

$$\bm{\mathcal{B}}^{e}=\bm{\mathcal{D}}^{e}\ \bm{W}^{e}$$

(29)

where $\bm{W}^{e}$ consists of the $J$ eigenvectors of the eigenvalue decomposition $\left(\bm{\mathcal{D}}^{e}\right)^{T}\bm{\mathcal{D}}^{e}=\bm{W}^{e}\bm{\Lambda}\left(\bm{W}^{e}\right)^{T}$.

Next, we apply $k$-means clustering method to partition the rows of the matrix $\bm{\mathcal{B}}^{e}$ into $K$ separate clusters. We denote the centroids of the clusters by $\bm{\mathcal{C}}^{e}_{k},1\leq k\leq K,1\leq e\leq N_{e}$. The clustering scheme allows us to divide the diverse dataset into smaller subsets, each characterized by similar data points which share the common attributes. This approach captures the diversity inherent in the dataset by identifying distinct atomic environments within the dataset. Figure 2 illustrates the process of partitioning the dataset into several atom clusters.

While local descriptors are used to partition $N$ atoms of the data set into $K$ subsets of atoms for each element, global descriptors can be used for partitioning the configurations of the dataset into $K$ subsets of configurations. For the dataset of $N_{\rm config}$ configurations, we sum the relevant rows of the local descriptor matrix $\bm{\mathcal{D}}$ to obtain the global descriptor matrix $\bm{\mathcal{G}}$ of size $N_{\rm config}$ by $M$. We then apply the above procedure to $\bm{\mathcal{G}}$ to obtain the desired configuration subsets. By training potentials on these subsets separately, we can construct MLIPs that are tailored to specific atomic environments.

We begin by introducing the atomic energies associated with the partitioned subsets

$$\mathcal{E}_{ik}(\bm{R},\bm{Z})=\sum_{m=1}^{M}c_{mk}\mathcal{D}_{im}(\bm{R},\bm{Z}),\quad 1\leq k\leq K,$$

(30)

where the coefficients $c_{mk}$ are fitted against the QM data. For simplicity of exposition, the same local descriptors are used to define the atomic energies, although we allow for different local descriptors to be used for each subset. To construct a single potential energy surface, we introduce a many-body many-potential expansion

$$E_{i}(\bm{R},\bm{Z})=\sum_{k=1}^{K}\mathcal{P}_{ik}(\bm{R},\bm{Z})\mathcal{E}_{ik}(\bm{R},\bm{Z})$$

(31)

where

$$\sum_{k=1}^{K}\mathcal{P}_{ik}(\bm{R},\bm{Z})=1,\quad\mathcal{P}_{ik}(\bm{R},\bm{Z})\geq 0.$$

(32)

Thus, the atomic energy at atom $i$ is a weighted sum of the individual contributions from the $K$ subsets. Note that $\mathcal{P}_{ik}$ denotes the probability of atom $i$ belonging to the $k$th subset. Like the local descriptors, the probabilities depend on the local atomic environment of the central atom $i$.

By inserting (30) into the many-potential expansion (31) and summing over index $i$, we obtain the PES as

$$E_{\rm T}(\bm{R},\bm{Z})=\sum_{i=1}^{N_{a}}\sum_{k=1}^{K}\sum_{m=1}^{M}c_{mk}\mathcal{P}_{ik}(\bm{R},\bm{Z})\mathcal{D}_{im}(\bm{R},\bm{Z}).$$

(33)

The quantities $\mathcal{Q}_{ikm}(\bm{R},\bm{Z})=\mathcal{P}_{ik}(\bm{R},\bm{Z})\mathcal{D}_{im}(\bm{R},\bm{Z})$ shall be called environment-adaptive descriptors. Hence, we can write the PES as follows

$$E_{\rm T}(\bm{R},\bm{Z})=\sum_{i=1}^{N_{a}}\sum_{l=1}^{KM}c_{l}\mathcal{Q}_{il}(\bm{R},\bm{Z}),$$

(34)

where $l$ is a linear indexing of $k$ and $m$. The coefficients $c_{l}$ are sought as solution of a least squares problem that minimizes a loss function defined as the weighted mean squared errors between the predicted energies/forces and the QM energies/forces for all configurations in the training dataset.

For the single cluster case $K=1$, the EA model (33) reduces to the standard linear model

$$E_{\rm T}(\bm{R},\bm{Z})=\sum_{i=1}^{N_{a}}\sum_{m=1}^{M}c_{m}\mathcal{D}_{im}(\bm{R},\bm{Z}).$$

(35)

Hence, the EA model (34) has $K$ times more descriptors and coefficients than the standard linear model (35). Because the size of the EA model increases with $K$, it can describe diverse environments in the dataset better than the standard linear model. Nonetheless, it is necessary to assess the EA model while varying $K$ and compare it against the standard linear model.

It remains to calculate the probabilities $\mathcal{P}_{ik}$. They are defined in terms of the local descriptors as follows. First, we compute the low-dimensional descriptors.

$${\mathcal{B}}_{ij}(\bm{R},\bm{Z})=\sum_{m=1}^{M}W_{mj}^{Z_{i}}{\mathcal{D}}_{im}(\bm{R},\bm{Z}),\quad 1\leq j\leq J.$$

(36)

where $W_{mj}^{e},1\leq e\leq N_{e},$ are the PCA matrices. Next, we calculate the inverse of the square of the distance from the $k$th centroid as

$${\mathcal{S}}_{ik}(\bm{R},\bm{Z})=\frac{1}{\sum_{j=1}^{J}\left({\mathcal{B}}_{ij}(\bm{R},\bm{Z})-\mathcal{C}_{kj}^{Z_{i}}\right)^{2}}.$$

(37)

Recall that $\mathcal{C}_{kj}^{e}$ are the centroids obtained from partitioning the dataset. Finally, the probabilities are calculated as

$${\mathcal{P}}_{ik}(\bm{R},\bm{Z})=\frac{{\mathcal{S}}_{ik}(\bm{R},\bm{Z})}{\sum_{l=1}^{K}{\mathcal{S}}_{il}(\bm{R},\bm{Z})},\quad 1\leq k\leq K.$$

(38)

Hence, the probabilities are high when the distances between the low-dimensional descriptor vector and the centroids are small. The additional cost of evaluating the probabilities is only $O(JM+JK)$ per atom. This cost can be neglected, as it is considerably less than the cost of computing the local descriptors.

Forces on atoms are calculated by differentiating the PES (33) with respect to atom positions. To this end, we first compute the partial derivatives of the probabilities with respect to the local descriptors as

$$\frac{\partial{\mathcal{P}}_{ik}}{\partial\mathcal{D}_{im}}=\frac{\partial{\mathcal{P}}_{ik}}{\partial\mathcal{S}_{il}}\frac{\partial{\mathcal{S}}_{il}}{\partial\mathcal{B}_{ij}}\frac{\partial{\mathcal{B}}_{ij}}{\partial\mathcal{D}_{im}}.$$

(39)

Here the Einstein summation convention is used to indicate the implicit summation over repeated indices except for the index $i$. The cost of evaluating the terms in (39) is $O(K^{2}MJ)$ per atom. Next, we note that

$$\frac{\partial{\mathcal{P}}_{ik}(\bm{R},\bm{Z})}{\partial\bm{R}}=\frac{\partial{\mathcal{P}}_{ik}}{\partial\mathcal{D}_{im}}\frac{\partial{\mathcal{D}}_{im}}{\partial\bm{R}}.$$

(40)

Differentiating the PES (33) with respect to atom positions yields

$$\frac{\partial E_{\rm T}(\bm{R},\bm{Z})}{\partial\bm{R}}=c_{nk}\mathcal{D}_{in}\frac{\partial{\mathcal{P}}_{ik}}{\partial\bm{R}}+c_{ml}\mathcal{P}_{il}\frac{\partial{\mathcal{D}}_{im}}{\partial\bm{R}}.$$

(41)

By inserting (40) into (41), we obtain

$$\frac{\partial E_{\rm T}(\bm{R},\bm{Z})}{\partial\bm{R}}=\left(c_{nk}\mathcal{D}_{in}\frac{\partial{\mathcal{P}}_{ik}}{\partial\mathcal{D}_{im}}+c_{ml}\mathcal{P}_{il}\right)\frac{\partial{\mathcal{D}}_{im}}{\partial\bm{R}}.$$

(42)

The additional cost of evaluating the terms in the parenthesis is only $O(3MK)$ per atom.