Fast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations

F ast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations Henrik Christiansen, Takashi Maruyama, Federico Errica, Viktor Zaverkin, Makoto Takamoto, and Francesco Alesiani NEC Laboratories Europe GmbH, Kurfürsten-Anlage 36, 69115 Heidelberg, Germany (Dated: March 27, 2025) We present an end-to-end differentiable molecular simulation framework (DIMOS) for molecular dynamics and Monte Carlo simulations. DIMOS easily integrates machine-learning-based inter-atomic potentials and implements classical force fields including particle-mesh Ewald electrostatics. Thanks to its modularity, both classical and machine-learning-based approaches can be easily combined into a hybrid description of the system (ML/MM). The superior performance and the high versatility is probed in different benchmarks and applications, with speed-up factors of up to 170 . 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 to ad-hoc chosen simulation parameters. Molecular simulations are a cornerstone of modern computational physics, chemistry and biology, enabling researchers to understand complex properties of the system [1]. Traditional molecular dynamics (MD) and Markov Chain Monte Carlo (MCMC) simulations rely on pre-defined force fields and specialized software to achieve large timescales and efficient sampling of rugged free-energy landscapes [2]. However, conventional MD and MCMC simulation packages generally lack the flexibility and modularity to easily incorporate cutting-edge computational techniques such as machine learning (ML) based enhancements: Advances in machine learning in-teratomic potentials (MLIPs) promise improved accuracy for MD simulations [3], yet integrating these techniques into a scalable and user-friendly framework remains a major challenge, especially when developing novel approaches [4]. Here we present an end-to-end differentiable molecular simulation framework (DIMOS) implemented in PyTorch [5], a popular library for ML research. DI-MOS implements essential algorithms to perform MD and MCMC simulations, providing an easy-to-use way to interface MLIPs and an efficient implementation of classical force field components in addition to implementations of common integrators and barostats. Additional components are the efficient calculation of neighborlists and constraint algorithms which allow for larger timesteps of the numerical integrator. By relying on PyTorch, we inherit many advances achieved by the ML community: We achieve fast execution speed on diverse hardware platforms, combined with a simple-to-use and modular interface implemented in Python.