Quantum-based Molecular Dynamics Simulations Using Tensor Cores

Tensor cores, along with tensor processing units, represent a new form of hardware acceleration specifically designed for deep neural network calculations in artificial intelligence applications. Tensor cores provide extraordinary computational speed and energy efficiency, but with the caveat that they were designed for tensor contractions (matrix-matrix multiplications) using only low-precision floating point operations. In spite of this, we demonstrate how tensor cores can be applied with high efficiency to the challenging and numerically sensitive problem of quantum-based Born-Oppenheimer molecular dynamics, which requires highly accurate electronic structure optimizations and conservative force evaluations. The interatomic forces are calculated on-the-fly from an electronic structure that is obtained from a generalized deep neural network, where the computational structure naturally takes advantage of the exceptional processing power of the tensor cores and allows for high performance in excess of 100 Tflops on the tensor cores of a single Nvidia A100 GPU. Stable molecular dynamics trajectories are generated using the framework of extended Lagrangian Born-Oppenheimer molecular dynamics, which combines computational efficiency with long-term stability, even when using approximate charge relaxations and force evaluations that are limited in accuracy by the numerically noisy conditions caused by the low precision tensor core floating-point operations. A canonical ensemble simulation scheme is also presented, where the additional numerical noise in the calculated forces is absorbed into a Langevin-like dynamics.

Simple and efficient algorithms for training machine learning potentials to force data

Abstract Machine learning models, trained on data from ab initio quantum simulations, are yielding molecular dynamics potentials with unprecedented accuracy. One limiting factor is the quantity of available training data, which can be expensive to obtain. A quantum simulation often provides all atomic forces, in addition to the total energy of the system. These forces provide much more information than the energy alone. It may appear that training a model to this large quantity of force data would introduce significant computational costs. Actually, training to all available force data should only be a few times more expensive than training to energies alone. Here, we present a new algorithm for efficient force training, and benchmark its accuracy by training to forces from real-world datasets for organic chemistry and bulk aluminum.

Automated discovery of a robust interatomic potential for aluminum

Accuracy of molecular dynamics simulations depends crucially on the interatomic potential used to generate forces. The gold standard would be first-principles quantum mechanics (QM) calculations, but these become prohibitively expensive at large simulation scales. Machine learning (ML) based potentials aim for faithful emulation of QM at drastically reduced computational cost. The accuracy and robustness of an ML potential is primarily limited by the quality and diversity of the training dataset. Using the principles of active learning (AL), we present a highly automated approach to dataset construction. The strategy is to use the ML potential under development to sample new atomic configurations and, whenever a configuration is reached for which the ML uncertainty is sufficiently large, collect new QM data. Here, we seek to push the limits of automation, removing as much expert knowledge from the AL process as possible. All sampling is performed using MD simulations starting from an initially disordered configuration, and undergoing non-equilibrium dynamics as driven by time-varying applied temperatures. We demonstrate this approach by building an ML potential for aluminum (ANI-Al). After many AL iterations, ANI-Al teaches itself to predict properties like the radial distribution function in melt, liquid-solid coexistence curve, and crystal properties such as defect energies and barriers. To demonstrate transferability, we perform a 1.3M atom shock simulation, and show that ANI-Al predictions agree very well with DFT calculations on local atomic environments sampled from the nonequilibrium dynamics. Interestingly, the configurations appearing in shock appear to have been well sampled in the AL training dataset, in a way that we illustrate visually.