RLEKF: An Optimizer for Deep Potential with Ab Initio Accuracy

Training Methods in NNMD Packages. Nearly all NNMD packages mentioned above use Adam and SGD as the training procedure for their ease of use in NN frameworks. One challenge of these methods is their time-consuming model training. For example, it usually takes more than 100 epochs to systematically train an individual DP model in the DeePMD-kit software.

DP. The key steps of the DP training are shown in Fig. 1(a). For each atom $i$, the physical symmetries such as translational, rotational, and permutational invariances are integrated into the descriptor $\mathcal{D}_{i}$ through a three-layer fully connected network named embedding net. Then $\mathcal{D}_{i}$ is trained via a three-layered fitting net.

1. 1.

   Every snapshot of the molecule system consists of each atom’s 3D Cartesian coordinates $\mathbf{r}_{i}=(x_{i},y_{i},z_{i})\in R^{3},i=1,2,...,N_{a}$ which is then translated into neighbor list of atom $i$, $\mathcal{R}_{i}=\{\mathbf{r}_{ij}\in\mathbb{R}^{3}|\mid\mathbf{r}_{j}-\mathbf{r}_{i}\mid<r_{c}\}$ as the input of DP NN. Then, we gather it into its smooth version $\tilde{\mathcal{R}}_{i}\in\mathbb{R}^{N_{m}\times 4}$, $(\tilde{\mathcal{R}}_{i})_{j}=s(|\mathbf{r}_{ij}|)(1,\mathbf{r}_{ij}/|\mathbf{r}_{ij}|)\in\mathbb{R}^{4}$, where $s(x)=1/x$ when $x<r_{cs}$, $s(x)=0$ when $x>r_{c}$, decaying smoothly between the two thresholds, $N_{m}$ is the maximum length of all neighbor lists, $j$ is the neighbor index of atom $i$.

2. 2.

   Define the embedding net $\mathcal{G}_{i}\in\mathbb{R}^{N_{m}\times M}$, $\mathcal{G}_{i}=\mathcal{G}(s(|\mathbf{r}_{i\cdot}|))$, where $M$ is called symmetry order, $\mathcal{G}=\mathcal{E}_{2}\circ\mathcal{E}_{1}\circ\mathcal{E}_{0}$, $\mathcal{E}_{l}(X)=X+\tanh\left(XW_{l}+\boldsymbol{1}\otimes w_{l}\right),l\in\{1,2\}$, $\mathcal{E}_{0}(\mathbf{x})=\tanh\left(\mathbf{x}\otimes W_{0}+\boldsymbol{1}\otimes w_{0}\right)$, $|\mathbf{r}_{i\cdot}|$ and $\boldsymbol{1}\in\mathbb{R}^{N_{m}\times 1},w_{l}\in\mathbb{R}^{1\times M},W_{l}\in\mathbb{R}^{N_{m}\times M},l\in\{0,1,2\}$ and the functions $s$ and $\tanh$ are element-wise.

3. 3.

   $\tilde{\mathcal{R}}_{i}$ timing $\mathcal{G}_{i}$ yields the so-called descriptor $\mathcal{D}_{i}:=\mathcal{G}_{i}^{\textsf{T}}\tilde{\mathcal{R}}_{i}\tilde{\mathcal{R}}_{i}^{\textsf{T}}\mathcal{G}_{i}^{<}\in\mathbb{R}^{M\times M^{<}}$, i.e. $\mathcal{G}_{i}^{<}$ is the several columns of $\mathcal{G}_{i}$ and $M^{<}<M$.

4. 4.

   The fitting net is $E_{i}=\mathcal{F}(\mathcal{D}_{i})=\mathcal{F}_{3}\circ\mathcal{F}_{2}\circ\mathcal{F}_{1}\circ\mathcal{F}_{0}(\mathcal{D}_{i})$, where $\mathcal{D}_{i}$ is reshaped into a vector of form $\mathbb{R}^{MM^{<}\times 1}$, $\mathcal{F}_{l}(\mathbf{x})=\mathbf{x}+\tanh\left(\tilde{W}_{l}\mathbf{x}+\tilde{w}_{l}\right),\tilde{w}_{l}\in\mathbb{R}^{d\times 1},\tilde{W}_{l}\in\mathbb{R}^{d\times d},l\in\{1,2\}$, $\mathcal{F}_{3}(\mathbf{x})=\tilde{W}_{3}\mathbf{x}+\tilde{w}_{3}$, $\tilde{w}_{3}\in\mathbb{R}$, $\tilde{W}_{3}\in\mathbb{R}^{1\times d}$, $\mathcal{F}_{0}(\mathbf{x})=\tanh\left(\tilde{W}_{0}\mathbf{x}+\tilde{w}_{0}\right)$ , $\tilde{w}_{0}\in\mathbb{R}^{d\times 1}$, $\tilde{W}_{0}\in\mathbb{R}^{d\times MM^{<}}$.

5. 5.

   The output $E:=\sum_{i}E_{i}$, $F_{i}:=-\nabla_{\mathbf{r}_{i}}E$.

EKF with Memory Factor for NNs. For neural networks, the model of interest

is formulated in stochastic language as an EKF problem targeting on $\boldsymbol{\hat{\theta}}_{t}$, where $\boldsymbol{w}$ is the vector of all trainable parameters in the network $h(\cdot,\cdot)$, $\{(x_{t},y_{t})\}_{t\in\mathbb{N}}$ are pairs of feature and label, $\{y_{t}\}_{t\in{\mathbb{N}}}$ can also be seen as measurements of EKF, $\{\eta_{t}\}_{t\in\mathbb{N}}$ are noise terms subject to normal distribution with mean $0$ and variances $\{R_{t}\}_{t\in\mathbb{N}}$ correspondingly, and $\forall t\in\mathbb{N},\boldsymbol{\hat{\theta}}_{t|t-1}:=\mathbb{E}[\boldsymbol{\theta}_{t}|y_{t-1},y_{k-2},\dots,y_{1}]$. With fixed $\boldsymbol{\theta}_{t}$ and bounded $x_{t}$, $h(\boldsymbol{\theta}_{t},x_{t})$ will be approximated well by its linearization at $\boldsymbol{\hat{\theta}}_{t|t-1}$, just omitting a term $\mathcal{O}((\boldsymbol{\theta}_{t}-\boldsymbol{\hat{\theta}}_{t|t-1})^{2})$.

If set $m_{t}=y_{t}-h(\boldsymbol{\hat{\theta}}_{t|t-1},x_{t})+\mathbf{H}_{t}^{\textsf{T}}\boldsymbol{\hat{\theta}}_{t|t-1}$ and rewrite (2) the following KF problem

At the beginning of training, the estimator $\boldsymbol{\hat{\theta}}_{t|t-1}$ is far away from $\boldsymbol{w}$, so less attention should be paid to those data fed to the network at an earlier stage of training than those at later stage. Through timing a factor $\alpha_{t}:=\Pi_{i=1}^{t}\lambda_{i}^{-1/2}$ and $\alpha_{0}:=1$, where $0<\lambda_{i}\leq 1$ and $\lambda_{i}\rightarrow 1$, the last problem enjoys the better variant as below

where $\lambda_{t}$ is called memory factor. The greater $\lambda_{t}$ is, the more weight, or say attention, is paid to previous data. According to basic KF theory (Haykin and Haykin 2001) , we obtain

Finally, we recover the estimator of $w$ via that of $\theta$ divided by the factor $\alpha_{t}$, define $\forall t\in\mathbb{N},\hat{\boldsymbol{w}}_{t\mid t-1}:=\alpha_{t}^{-1}\hat{\boldsymbol{\theta}}_{t\mid t-1},\boldsymbol{w}_{t}:=\alpha_{t}^{-1}\hat{\boldsymbol{\theta}}_{t}$, find $\hat{\boldsymbol{w}}_{t\mid t-1}=\boldsymbol{w}_{t-1}$, and then get our weights updating strategy

The Workflow of Training with RLEKF. The overview of DP NN with RLEKF is shown in Fig. 1, and the corresponding training procedures are detailed in Alg. 1. Specifically, Alg. 1 and Fig. 1.b show the training of RLEKF for both energy (Alg. 1, line 2 & 3 ) and force (Alg. 1, line 4 & 5 ). The forward prediction(shown in Alg. 1 and Fig. 1.c), backward propagation (shown in Alg. 2 and Fig. 1.c), and updating of weights based on Kalman filtering with forgetting factor are shown in Alg. 2 (line 8 to line 14 and Fig. 1.e). Alg. 1 shows the iterative update of $\boldsymbol{w}_{t}$ with $\mathbf{P}_{t},\lambda_{t}$ (also shown in Fig. 1.b). After the initialization of $\boldsymbol{w}_{0}$, $\mathbf{P}_{0}$, and $\lambda_{1}$, the energy process starts and tries to fit the predicted energy to the energy label of a sample, and then the force process updates weights under the supervision of the samples’ force label. In the force process for each image, we randomly choose $x$ atoms and concatenate their network-predicted force vectors as an effective predicted force $\hat{F}$, which is then used to update weights. This process (i.e. line 4 & 5 of Alg. 1) is repeated $y$ times in a loop. Here, we recommend a relatively universal setting $x=6,y=4$. When these two kinds of alternative weights update in the direction of minimizing the trace of $\mathbf{P}_{t}$ for this sample, the same process for the next one repeats until time step $T$. The algorithm RLEKF in line 3 & 5 of Alg. 1 unfolds below.

RLEKF. Alg. 2 shows the layerwise weights-updating strategy of RLEKF, which is an optimized version of EKF in training the DP model. From line 1 to line 7, we transform any vector $\hat{Y}$ into a number, which enjoys two advantages. The first is it averages all the updating information in $Y^{\text{DFT}}_{i}-\hat{Y}_{i}$ and transforms a vector into a number which avoids an impregnable problem, solving inversion of overwhelmingly many big matrices introduced later. The second is once Alg. 2 invoked we only need to calculate the gradient for one time. More specifically, from line 1 to line 7, we just want to adjust the gradient of $\hat{Y}_{i}$ with respect to weights in the direction of ”decreasing the difference between $\hat{Y}_{i}$ and $Y^{\text{DFT}}_{i}$”. In order to keep $\mathbf{P}^{out}$ strictly symmetric if $\mathbf{P}^{in}$ is, the definition of Kalman gain in Alg. 2 is a little different from traditional Kalman filtering. In addition, $\mathbf{a}$ is usually a matrix, but in line 8 it degenerates into a number, which heavily reduces the computational complexity, where $L$ is the number of blocks in $\mathbf{P}^{in}$ and take $\alpha_{t}^{2}R_{t}=L\mathbf{I}$ as a suitable hyperparameter. RLEKF independently updates weights and $\mathbf{P}_{t}$ of different layers by KF theory and its splitting strategy introduced in the next subsection. Finally, all updated parameters are collected for the prediction of the next turn (Fig. 1.e, Alg. 2 from line 8 to line 16, where $\nu$ is forgetting rate, a hyperparameter describing the varying rate of $\lambda_{t}$ and split is a function for obtaining the $l$-th part after splitting).

The Splitting Strategy of RLEKF. According to the splitting strategy of RLEKF, GEKF can be approximated by RLEKF with a reduction in weights error covariance matrix considering the efficiency and stability of a large-scale application. Unlike GEKF, the weights error covariance matrix of whichever time step, up to a scalar factor, $\mathbf{P}=\left\{P_{1},\dots,P_{L}\right\}$ is a $N\times N$ block diagonal matrix, whose shape is $\left\{n_{1}\times n_{1},n_{2}\times n_{2},\dots,n_{L}\times n_{L}\right\}$, where $n_{i}$ and $N:=\sum_{i}n_{i}$ are the number of weights of the $i$th block and that of the whole NN respectively. This means some weights error covariances are forced into $0$ and $\mathbf{P}$ thus is no longer the real weights error covariance matrix at most an approximation, but fortunately there are indeed such good enough approximations that they carry most of the ”big values” and information in these diagonal blocks, on which $\mathbf{P}$ concentrates. As shown in Fig. 2, GEKF concerns correlations between all the parameters. However, it is computationally expensive when adapted to large NNs. As a seemingly good choice, SEKF is a load balanced approximation version of GEKF, but the heavy correlations between parameters in a layer are ignored inappropriately. Overcoming the drawback of SEKF, LEKF can fully consider the internal correlation between parameters in the same layer, whereas unbearable computation still confronts us if the layer contains tens of thousands of parameters and decoupling every small neighboring layer deviates from the load balance idea. Here, we induce two heuristic principles for choosing a suitable strategy:

1. 1.

   Put weights with high correlation (in the same layer) in the same block if possible.

2. 2.

   The block must not be too large since that burdens computers with much computation (splitting the layers if the number of parameters is larger than the threshold $N_{b}$) or too small since that wastes computational power and loses massive information (gathering the near neighboring small layers). The threshold $N_{b}$ aims at splitting the matrix $\mathbf{P}$ as evenly as possible for load balance and linear computational complexity.

Therefore, following these principles, we induce the splitting strategy of RLEKF (Fig. 1.d) and reorganize the weights in different layers, named parameter parts, into several more appropriate layers. Trainable parameters could be decomposed into a series of $l$ the parts of size $p_{1},p_{2},\dots,p_{l}$. Starting from the first part of size $p_{1}$, we execute the commands below repeatedly until $z=l$:

1. 1.

   For a given part with $p_{s}$ weights, we split them into $\lfloor\frac{p_{s}}{N_{b}}\rfloor$ layers of size $N_{b}$ and a new part of size $p_{s}\leftarrow p_{s}-N_{b}\lfloor\frac{p_{s}}{N_{b}}\rfloor$.

2. 2.

   Then consider subsequent parts with $p_{s},p_{s+1},\dots,p_{z}$ weights, gather them together into a layer of size $\sum_{i=s}^{z}p_{i}$ and then set $s\leftarrow z+1$ , if $\sum_{i=s}^{z}p_{i}\leq N_{b}$ and $(\sum_{i=s}^{z+1}p_{i}>N_{b}$ or $z=l)$.

Empirically, Fig. 2 validates the efficiency of RLEKF, comparably or even more accurate than GEKF.

Experiment Setting. Set $\lambda_{1}=0.98$, $\mathbf{P}_{0}=\mathbf{I}$, $\nu=0.9987$, $\boldsymbol{w}_{0}$ consistent with DeePMD-kit (Wang and Weinan 2018). The network configuration is [1, 25, 25, 25] (embedding net), [400, 50, 50, 50,1] (fitting net).

Data Description. We choose 13 unbiased and representative datasets of the aforementioned systems with certain specific structures (second column of Tab. 1). For each dataset, snapshots are yielded based on solving ab initio molecular trajectories via PWmat (Jia et al. 2013). During this process, to enlarge the sampling span of configuration space, we fast generate a long sequence of the snapshot by small time step (the third column of Tab. 1) and choose one for every fixed number in the temperature 300K.

Main Results. Both RLEKF and Adam yield good results except for CuO and Cu$+$C systems (Tab. 1). From an accuracy perspective, RLEKF reaches a higher energy accuracy in 12 cases except for Si, little less accurate than Adam. As for force, RLEKF is more accurate than Adam in 7 cases while the reminder achieves a very comparable precision. Furthermore, Fig. 3 shows how far the fitting results of RLEKF deviate from those of DFT (ground truth). From a speed perspective, generally, RLEKF converges to a reasonable RMSE much faster than Adam in most of the 13 cases in terms of energy (Fig. 4). We take bulk copper as an example to demonstrate how to understand information in Fig. 4. The lower Energy and Force RMSE between Adam and RLEKF is 0.327 meV and 44.4 meV/$\mathring{A}$. Therefore, the 1.2$\times$ energy baseline is 0.327$\times$1.2=0.392 meV. It is unreachable for Adam, which is denoted as 1000 in (Fig. 5). The 1.2$\times$ force baseline is 44.4$\times$1.2=52.6meV/$\mathring{\text{A}}$. Adam spends 35 epochs to reach the baseline while RLEKF’s only costs 3 epochs (Fig. 5). Then, the force speed ratio is 35/3=11.67 according to epoch. For each sample, Adam updates weights for 1 time and RLEKF updates for 7 times (1 time in the energy process and 6 times in the force process). The iteration speed ratio is 35/(3$\times$7)=1.66. On tesla V100, the time cost of each epoch is 60s (Adam) and 587s (RLEKF). Thus, the wall-clock time speed ratio is 35$\times$60/(3$\times$587)=1.19.

Robustness Analysis: Distribution of Hyperparameter of Weights Initialization. RLEKF is very stable as an optimizer, which can keep NNs from gradients exploding and consequently endow them with very loose initialization constraints almost up to none as shown in Fig. 6.

Computational Complexity. RLEKF also benefits from computing through reducing the computation compared to GEKF. There are 3 computational intensive parts in Alg. 2, calculating $a$ (line 8), $K$ (line 10), and $P^{out}$ (line 11). Due to the even splitting strategy, the order of float operation for each block is $\mathcal{O}(N_{b}^{2})$, and the number of the block is of order $\mathcal{O}(N/N_{b})$. Hence the total computational complexity of RLEKF is of order $\mathcal{O}((N/N_{b})N_{b}^{2})=\mathcal{O}(NN_{b})$.