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

03 January 2022, Version 2
This content is a preprint and has not undergone peer review at the time of posting.


Molecular dynamics (MD) simulations facilitate the study of physical and chemical processes of interest. Traditional classical MD models lack reactivity to explore several important phenomena; while quantum mechanical (QM) models can be used for this purpose, they come with steep computational costs. The reactive force field (ReaxFF) model bridges the gap between these approaches by incorporating dynamic bonding and polarizability. To achieve realistic simulations using ReaxFF, model parameters must be optimized against high fidelity training data, typically with QM accuracy. 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 can require millions of error evaluations for complex parameter fitting tasks, significantly hampering the rapid development of high quality parameter sets. In this work, we present JAX ReaxFF, a novel software tool that leverages modern machine learning infrastructure to enable extremely fast optimization of ReaxFF parameters. By calculating gradients of the loss function using the JAX library, we are able to utilize highly effective local optimization methods, such as the limited Broyden–Fletcher–Goldfarb–Shanno (LBFGS) and Sequential Least Squares Programming (SLSQP) methods. As a result of the performance portability of JAX, JAX-ReaxFF can execute efficiently on multi-core CPUs, GPUs (or even TPUs). By leveraging the gradient information and modern hardware accelerators, we are able to decrease parameter optimization time for ReaxFF from days to mere minutes. JAX-ReaxFF framework can also serve as a sandbox environment for domain scientists to explore customizing the ReaxFF functional form for more accurate modeling.


force field optimization

Supplementary weblinks


Comments are not moderated before they are posted, but they can be removed by the site moderators if they are found to be in contravention of our Commenting Policy [opens in a new tab] - please read this policy before you post. Comments should be used for scholarly discussion of the content in question. You can find more information about how to use the commenting feature here [opens in a new tab] .
This site is protected by reCAPTCHA and the Google Privacy Policy [opens in a new tab] and Terms of Service [opens in a new tab] apply.