LEARNING FORCE FIELD PARAMETERS FROM DIFFERENTIABLE MOLECULAR DYNAMICS

We developed Diff-MD, a fully end-to-end differentiable molecular dynamics software, and use it to automatically derive force field (FF) parameters. Diff-MD is templated on the recently released ∂-HylleraasMD software and leverages an extensive list of python libraries developed for machine learning. The differentiable simulation is implemented such that every calculation, from determining interatomic distances to updating positions, can be automatically differentiated. This enables the computation of gradients for any property that depends on the simulation trajectory (e.g., density, area per lipid, polymer chain conformation) with respect to the FF parameters. With these gradients, we employ gradient-based optimization algorithms to refine the FF parameters. Here, we used Diff-MD to derive the interaction energy parameters, ε, of the Lennard-Jones potential for coarse grained models of phospholipid membranes and polymers. As target properties, we selected the area per lipid and density profiles for membranes, and the radius of gyration and end-to-end distance for polymer chains. By using massive parallelization, we optimize a set of systems of different compositions and in different conditions, ensuring the transferability of the parameters obtained. We use the JAX library to obtain gradients, while implementing the force calculation, MD integrator, barostat, thermostat and neighbor list generation using the JAX NumPy API, and taking advantage of Just In Time compilations whenever possible to accelerate the calculations. The parameters are restrained to a physical range of values using the Optax library. Diff-MD can incorporate target properties from both atomistic simulations and experiments, and its framework allows for the inclusion of additional properties beyond those explored here.