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

H Christiansen and T Maruyama and F Errica and V Zaverkin and M Takamoto and F Alesiani, JOURNAL OF CHEMICAL PHYSICS, 163, 182501 (2025).

DOI: 10.1063/5.0277356

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 170x 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 3x acceleration is observed in comparison with ad-hoc chosen simulation parameters. The code is available at https://github.com/nec- research/DIMOS.

Return to Publications page