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

arXiv: https://arxiv.org/abs/2503.20541 · HTML: https://mlsanity.com/paper/2503.20541/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 60 references that are in the map.
- First posted on arXiv 26 Mar 2025.

## Builds on (most cited first, 13 of 13)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems](https://mlsanity.com/paper/2312.07511/) (2023)
- [Orb: A Fast, Scalable Neural Network Potential](https://mlsanity.com/paper/2410.22570/) (2024)
- [TorchMD: A deep learning framework for molecular simulations](https://mlsanity.com/paper/2012.12106/) (2020)
- [OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials](https://mlsanity.com/paper/2310.03121/) (2023)
- [Adaptive Message Passing: A General Framework to Mitigate Oversmoothing, Oversquashing, and Underreaching](https://mlsanity.com/paper/2312.16560/) (2023)
- [Higher-Rank Irreducible Cartesian Tensors for Equivariant Message Passing](https://mlsanity.com/paper/2405.14253/) (2024)
- [NNP/MM: Accelerating molecular dynamics simulations with machine learning potentials and molecular mechanic](https://mlsanity.com/paper/2201.08110/) (2022)
- [FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials](https://mlsanity.com/paper/2405.01491/) (2024)
- [Self-Tuning Hamiltonian Monte Carlo for Accelerated Sampling](https://mlsanity.com/paper/2309.13593/) (2023)
- [Adaptive Width Neural Networks](https://mlsanity.com/paper/2501.15889/) (2025)

## Built on by (most cited first, 2 of 2)

- [Performance of universal machine-learned potentials with explicit long-range interactions in biomolecular simulations](https://mlsanity.com/paper/2508.10841/) (2025)
- [Enabling Structure-Only Initialization and Out-of-Distribution Generalization in GNN-based Molecular Dynamics Simulators](https://mlsanity.com/paper/2605.09495/) (2026)

## Abstract (arXiv)

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 (ML/MM). 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 verb|PyTorch| implementation. We show that due to improved linear instead of quadratic scaling as 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 (HMC). Using these optimized simulation parameters a 3× acceleration is observed in comparison to ad-hoc chosen simulation parameters. The code is available at https://github.com/nec-research/DIMOS.