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 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). By supporting key molecular dynamics features such as efficient neighborlists and constraint algorithms for larger time steps, the framework bridges the gap between hand-optimized simulation engines and the flexibility of a PyTorch implementation. The superior performance and the high versatility is probed in different benchmarks and applications, with speed-up factors of up to $170\times$. 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\times$ acceleration is observed in comparison to ad-hoc chosen simulation parameters. The code is available atthis https URL.

@article{christiansen2025_2503.20541,
  title={ Fast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations },
  author={ Henrik Christiansen and Takashi Maruyama and Federico Errica and Viktor Zaverkin and Makoto Takamoto and Francesco Alesiani },
  journal={arXiv preprint arXiv:2503.20541},
  year={ 2025 }
}