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 (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\times$ 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\times$ acceleration is observed in comparison to ad-hoc chosen simulation parameters. The code is available at https://github.com/nec-research/DIMOS.

💡 Research Summary

The paper “Fast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations” introduces DIMOS, a novel end-to-end differentiable molecular simulation framework built on PyTorch. It aims to bridge the gap between the high performance of traditional, hand-optimized simulation engines (like LAMMPS) and the flexibility and differentiability required for modern machine learning (ML) research in computational chemistry and materials science.

DIMOS supports both Molecular Dynamics (MD) and Markov Chain Monte Carlo (MCMC) simulations. Its core design principle is modularity, allowing seamless integration of classical molecular mechanics (MM) force fields and machine learning interatomic potentials (MLIPs), including hybrid ML/MM approaches. The framework implements essential high-performance MD algorithms within the PyTorch computational graph. Key technical features include: 1) Efficient bonded interactions (bonds, angles, dihedrals) following AMBER/CHARMM formats, 2) Non-bonded interactions handled via a combined Verlet neighbor list and cell list algorithm, achieving O(N) scaling with system size, 3) Electrostatics treated with either a reaction field approximation or a Particle-Mesh Ewald (PME) method for long-range interactions, and 4) Integrators (e.g., Langevin dynamics) and constraint algorithms for enabling larger timesteps.

The authors demonstrate significant performance advantages. Benchmarks on water boxes show that DIMOS achieves linear scaling, leading to speed-ups of up to 170x compared to another fully differentiable framework (torchMD) for classical force field simulations, primarily due to the efficient neighbor list implementation. The feasibility of running hybrid ML/MM simulations on protein systems is also shown.

A major innovation highlighted is the practical application of differentiability beyond just training MLIPs. The authors perform an end-to-end optimization of simulation meta-parameters. Using Hamiltonian Monte Carlo (HMC) for alanine dipeptide in explicit solvent, they define a loss function based on the difference between the sampled and target distributions and use automatic differentiation to optimize the HMC proposal distribution parameters (e.g., the mass matrix). This results in a 3x acceleration in sampling efficiency compared to using ad-hoc parameters, showcasing a novel paradigm for tuning simulation algorithms themselves.

In conclusion, DIMOS is positioned not as a direct replacement for heavily optimized production MD codes but as a complementary tool for rapid prototyping, method development, and research that leverages differentiability. It fills a niche in the existing ecosystem by providing a performant, modular, and fully differentiable simulation environment within the popular PyTorch ecosystem, opening doors for new research in force field optimization, enhanced sampling, and integrative ML/MM modeling. The code is publicly available on GitHub.