Learning Force Field Parameters from Differentiable Particle-Field Molecular Dynamics

We develop ∂-HylleraasMD (∂-HyMD), a fully end-to-end differentiable molecular dynamics software based on the Hamiltonian hybrid particle-field formalism, and use it to establish a protocol for automated optimization of force field parameters. ∂-HyMD is templated on the recently released HylleraaasMD software, while using the JAX autodiff framework as the main engine for the differentiable dynamics. ∂-HyMD exploits an embarrassingly parallel optimization algorithm by spawning independent simulations, whose trajectories are simultaneously processed by reverse mode automatic differentiation to calculate the gradient of the loss function, which is in turn used for iterative optimization of the force-field parameters. We show that parallel organization facilitates the convergence of the minimization procedure, avoiding the known memory and numerical stability issues of differentiable molecular dynamics approaches. We showcase the effectiveness of our implementation by producing a library of force field parameters for standard phospholipids, with either zwitterionic or anionic heads and with saturated or unsaturated tails. Compared to the all-atom reference, the force field obtained by ∂-HyMD yields better density profiles than the parameters derived from previously utilized gradient-free optimization procedures. Moreover, ∂-HyMD models can predict with good accuracy properties not included in the learning objective, such as lateral pressure profiles, and are transferable to other systems, including triglycerides.