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

TL;DR

DIMOS provides a PyTorch-based, end-to-end differentiable framework for molecular dynamics and Monte Carlo simulations, unifying classical force fields with machine-learning interatomic potentials (MLIPs) and ML/MM hybrids. It achieves favorable scaling through neighborlists, PME/PME-like long-range electrostatics, and GPU-accelerated differentiable computations, offering substantial speedups over existing differentiable engines. The paper demonstrates practical benefits across water boxes and protein systems, including large ML/MM speedups, and showcases end-to-end differentiability by optimizing HMC parameters to significantly reduce sampling autocorrelation. Overall, DIMOS serves as a flexible, gradient-enabled platform for rapid prototyping and method development in computational chemistry and biophysics, complementing but not replacing production MD codes.

Abstract

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.

Figures (4)

• Figure 1: In (a)-(c) a schematic representation of Alanine tripeptide (Ala$_3$), where the atomic interactions are modeled (a) solely classically (MM), (b) solely using an MLIP (ML), or (c) using a hybrid approach (ML/MM). The red circle marks a region modeled by a classical force field, where one defines interaction terms based on covalent bonds and non-covalently bonded atoms. The blue circles represent the message passing iterations of a graph neural network used to construct many-body interactions, for details see Ref. duval2023hitchhiker and the references therein. In all cases, only the operations based on calculating the force/energy for a single atom are depicted, and for a full evaluation the steps are repeated for every atom. In (d), the energy components for a classical force field are depicted, based on the parameterization of AMBER class interactions.weiner1981amber Covalently bonded and non-bonded interactions are distinguished, with exclusions in the non-bonded interactions for directly or indirectly bonded atoms. For more details on the individual contributions, see Section \ref{['sec:FF']}.
• Figure 2: Runtime $\tau$ in seconds per MD step for a water box with constant density as a function of number of atoms $N$. (a) Shows the runtimes for torchMD, DIMOS, and openMM using single precision numerics and RFA to model the electrostatic interactions. (b) Demonstrates the influence of different approaches to model the electrostatic interactions: Runtimes are presented for RFA, Ewald, and PME for DIMOS using single precision numerical accuracy. In (c), we investigate the effect of the chosen computational precision on the runtime, modelling electrostatic interactions via RFA in DIMOS. In all cases, the lines indicate the theoretically expected asymptotic computational complexity. (d)-(f) Show the performance of running the MACE foundation model, where in (d) we compare the runtimes of DIMOS to TorchSim, ASE, and LAMMPS using the small MACE model and mixed precision (apart from TorchSim, which uses single precision) for different $N$, and (e) investigates the influence of choosing different numerical precision within DIMOS. (f) Presents the runtimes for the biggest system that can be simulated using the large MACE model ($N=8844$) and compares the resulting $\tau$ for DIMOS and ASE.
• Figure 3: (a) Demonstration on how to set up both classically and machine-learning enhanced systems in DIMOS. specifies which atoms should be modeled by the MLIP system, and also can be empty in case one does not want to model any systems using ML. (b-d) Achievable simulation time per day $\mathcal{T}$ with and without constrained motion of the hydrogens for each of the three cases of a system modeled purely by ML, the hybrid approach using ML/MM, and a purely MM based modelling. In all cases, the MLIP model was chosen to be . In (b), we present the data for alanine-dipeptide (Ala$_2$), (c) shows Trp-Cage, and (d) shows dihydrofolate reductase (DHFR). All systems are modeled using explicit water modeled by the TIP3P parameters.jorgensen1983comparison In (d), the biggest system, it is not possible to model purely with ML, as for this system size the MLIP model required more than the available GPU memory.
• Figure 4: Plots demonstrating the end-to-end learning of simulation parameters for HMC. (a) Loss $\mathcal{L}_{\mathrm{HMC}}$ as function of training time $t$. (b),(c) The mean of timestep $\mu_{\Delta t}$ and number of integration steps $\mu_L$ against training time $t$. (d),(e) The behavior of the masses $m$ for (d) the hydrogen and oxygen atoms in the water molecules, optimized as shared parameter between all water molecules and (e) the remaining masses in Ala$_2$. In (e) we plot $m-m_0$ to better highlight the relative change of masses for the diverse atom types. (f) Autocorrelation function $A(k)$ of the potential energy $U(x)$ as function of lag-time $k$ measured in terms of the number of energy evaluations.