Fast, modular, and differentiable framework for machine learning-enhanced molecular simulations

We present an end-to-end differentiable molecular simulation framework (DIMOS) for molecular dynamics and Monte Carlo simulations. DIMOS easily integrates machine-learning-based interatomic potentials and implements classical force fields including an efficient implementation of particle-mesh Ewald. Thanks to its modularity, both classical and machine-learning-based approaches can be easily combined into a hybrid description of the system (machine learning/mechanics modeling). By supporting key molecular dynamics features, such as efficient neighborlists and constraint algorithms for larger time steps, the framework makes steps in bridging the gap between hand-optimized simulation engines and the flexibility of a PyTorch implementation. We show that due to improved linear scaling instead of quadratic scaling as a function of system size, DIMOS is able to obtain speed-up factors of up to 170× for classical force field simulations against another fully differentiable simulation framework. The advantage of differentiability is demonstrated by an end-to-end optimization of the proposal distribution in a Markov Chain Monte Carlo simulation based on Hamiltonian Monte Carlo. Using these optimized simulation parameters, a 3× acceleration is observed in comparison with ad-hoc chosen simulation parameters. The code is available at https://github.com/nec-research/DIMOS.