Active Learning of Neural Network Interatomic Potentials with Differentiable Uncertainty

Neural networks (NNs) are extremely effective interpolators in atomistic simulations. Given abundant and diverse data, NN interatomic potentials (NNIPs) can be trained to replicate potential energy surfaces with the accuracy of high-level and high-cost electronic structure methods while reducing the computational cost by orders of magnitude. However, NNIPs are notoriously delicate and struggle to generalize to points outside the training data, which may result in highly erroneous predictions for atomic configurations not seen at train time. Naturally, the purpose of NNIP is to perform simulations at much larger time and length scales than can be accessed with the ground truth method, meaning that rare events not represented in the training data are likely to be seen in production.<br/> <br/>Uncertainty estimates then become key to building self-correcting NNIPs. By assigning confidence to their predictions during production simulations, the NNIPs can signal that the model is uncertain and needs to be re-trained with new data that is representative of the new environment, in so-called active learning strategies. However, this still requires (a cycle of) large-scale simulations to encounter these rare environments, identify new training data, and restart the process.<br/> <br/>Here, we describe how NN uncertainty quantification methods based using both ensembles and Maximum Likelihood Estimation, enable gradient-based active learning to systematically improve NN potentials. Because the uncertainty metrics are differentiable end-to-end with respect to model inputs, and directly distort input atomic positions towards regions of high thermodynamic likelihood (low energy) and high uncertainty through gradient ascent methods. The geometries found are thus representative of high-uncertainty configurations, but they are visited through fast optimization methods, rather than waiting for a traditional simulation to visit them.<br/><br/>Here, we will describe how uncertainty-based adversarial attacks on uncertainty help train more accurate potentials, with fewer data points, in fewer iterations, and at lower computational cost, for a variety of systems ranging from a double-well to organic molecules in silica nanopores. Furthermore, differentiable uncertainty allows the application of attribution techniques to the uncertainty itself, thus highlighting the atomic configurations responsible for model uncertainty. By extracting only the uncertain region for calculation with the ground-truth method, the computational cost is reduced, since no compute time is wasted on well-understood regions of the simulation box.