Adaptive Loss Weighting for Machine Learning Interatomic Potentials

For atomistic simulations, machine learning interatomic potentials (ML-IAPs) have proven to accurately represent potential energy surfaces, overcoming some limitations of empirical force-fields and their functional form. Training such potentials involves optimization of multipart loss functions typically composed of potential energies, forces and stress tensors. However, the contribution of each variable to the total loss is typically weighted using heuristic approaches that yield either iterative or sub-optimal results. Therefore, we implement an adaptive loss weighting algorithm based on a mathematically intuitive and computationally efficient scheme, in which the contribution of each term is dynamically recalculated based on its learning performance; optimizing the imbalance of the widely used, heuristic fixed loss weighting approaches. Additionally, by adding a stress term to the loss function and using high convergence criteria for density functional theory (DFT) calculations we show that a ML-IAP with accurate predictive stress capabilities must include stress information during training. This approach results in more automated neural networks networks that can learn effectively potential energies, forces and stresses from <i>ab initio</i> calculations, producing reliable ML-IAPs that are usable in simulations involving mechanical deformations, phase transformations, or other stress-dependent phenomena.