RLEKF: An Optimizer for Deep Potential with Ab Initio Accuracy

TL;DR

This paper introduces RLEKF, a novel optimizer for deep potential models that accelerates training with ab initio accuracy, outperforming Adam in convergence speed and stability for large neural networks.

Contribution

We propose RLEKF, an optimized extended Kalman filtering method with a layer-splitting strategy, improving training efficiency and stability for deep potential neural networks.

Findings

RLEKF converges faster than Adam in training deep potential models.

RLEKF achieves slightly better accuracy than Adam.

Theoretical proof shows RLEKF updates are stable and prevent gradient explosion.

Abstract

It is imperative to accelerate the training of neural network force field such as Deep Potential, which usually requires thousands of images based on first-principles calculation and a couple of days to generate an accurate potential energy surface. To this end, we propose a novel optimizer named reorganized layer extended Kalman filtering (RLEKF), an optimized version of global extended Kalman filtering (GEKF) with a strategy of splitting big and gathering small layers to overcome the $O(N^2)$ computational cost of GEKF. This strategy provides an approximation of the dense weights error covariance matrix with a sparse diagonal block matrix for GEKF. We implement both RLEKF and the baseline Adam in our $\alpha$Dynamics package and numerical experiments are performed on 13 unbiased datasets. Overall, RLEKF converges faster with slightly better accuracy. For example, a test on a typical system, bulk copper, shows that RLEKF converges faster by both the number of training epochs ($\times$11.67) and wall-clock time ($\times$1.19). Besides, we theoretically prove that the updates of weights converge and thus are against the gradient exploding problem. Experimental results verify that RLEKF is not sensitive to the initialization of weights. The RLEKF sheds light on other AI-for-science applications where training a large neural network (with tons of thousands parameters) is a bottleneck.