Forcefield Parameterization via Differentiable Molecular Dynamics

In classical molecular dynamics, the quality of the forcefield is determined by how well it reproduces the desired physical properties in comparison to those achieved in experiment/<i>ab initio </i>settings. To achieve this the forcefield parameters are chosen either by physical intuition or are optimized with respect to some target (ground truth) properties like energies or forces. This makes it important to parametrize the forcefield in an accurate and efficient manner. Optimization of a forcefield is typically done using either gradient-free or gradient-based algorithms. In the case of gradient-based algorithms, the gradients of the loss function have to be computed. Since the analytic gradients are not known, approximations like the finite difference are made to compute the gradients. Such approximations introduce errors and their computation is inefficient in higher dimensions. Automated (algorithmic) differentiation (AD) solves this problem by enabling the computation of gradients in an efficient way. AD libraries like JAX have enabled end-to-end differentiable simulation packages like JAX-MD. JAX-MD provides a framework for physics simulation primarily focused on molecular dynamics. Differentiable simulation enabled by JAX and JAX-MD allows for the computation of gradients that are utilized for computing the relevant properties (forces etc) as well as in optimizers that use gradients of the loss function. This enables parametrization based on static (energy, forces, elastic constants, etc.) and dynamic properties (radial distribution function, etc.). We demonstrate this by parameterizing forcefields in an efficient manner by using gradient information provided by JAX and JAX-MD. We also compare it with gradient-free optimization methods and highlight the advantages offered by differentiable simulation.