JAX-ReaxFF: A Gradient-Based Framework for Fast Optimization of Reactive Force Fields

The reactive force field (ReaxFF) model bridges the gap between traditional classical models and quantum mechanical (QM) models by incorporating dynamic bonding and polarizability. To achieve realistic simulations using ReaxFF, model parameters must be optimized against high fidelity training data which typically come from QM calculations. Existing parameter optimization methods for ReaxFF consist of black box techniques using genetic algorithms or Monte Carlo methods. Due to the stochastic behavior of these methods, the optimization process oftentimes requires millions of error evaluations for complex parameter fitting tasks, thereby significantly hampering the rapid development of high quality parameter sets. Rapid optimization of the parameters is essential for developing and refining Reax force fields because producing a force field which exhibits empirical accuracy in terms of dynamics typically requires multiple refinements to the training data as well as to the parameters under optimization. In this work, we present JAX-ReaxFF, a novel software tool that leverages modern machine learning infrastructure to enable fast optimization of ReaxFF parameters. By calculating gradients of the loss function using the JAX library, JAX-ReaxFF utilizes highly effective local optimization methods that are initiated from multiple guesses in the high dimensional optimization space to obtain high quality results. Leveraging the architectural portability of the JAX framework, JAX-ReaxFF can execute efficiently on multicore CPUs, graphics processing units (GPUs), or even tensor processing units (TPUs). As a result of using the gradient information and modern hardware accelerators, we are able to decrease ReaxFF parameter optimization time from days to mere minutes. Furthermore, the JAX-ReaxFF framework can also serve as a sandbox environment for domain scientists to explore customizing the ReaxFF functional form for more accurate modeling.