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

2025-03-26

AI Summary

Traditional molecular simulations suffer from low computational efficiency, limited flexibility, and difficulties in jointly optimizing machine-learned potentials with classical force fields. To address these challenges, this paper introduces DIMOS, an end-to-end differentiable molecular simulation framework. DIMOS features a modular architecture that unifies molecular dynamics and Monte Carlo sampling, and—crucially—enables end-to-end differentiable optimization of proposal distributions in Hamiltonian Monte Carlo for the first time. It seamlessly integrates machine-learned potentials, classical force fields (including Lennard-Jones and particle-mesh Ewald electrostatics), efficient neighbor lists, RATTLE constraints, and automatic differentiation, while supporting ML/MM hybrid modeling. Experiments demonstrate up to 170× speedup over conventional simulation engines; robust and efficient ML/MM simulations; and a 3× improvement in sampling efficiency via joint optimization of HMC parameters.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Multimodal Learning
• Cognitive Modeling & Cognitive Systems: Simulating Human Behavior

Application Category

• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Web Mining and Content Analysis: Web data generation and simulation

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

Problem

• Develops differentiable molecular simulation framework (DIMOS) for dynamics and Monte Carlo
• Integrates machine-learning potentials with classical force fields (ML/MM hybrid)
• Optimizes simulation parameters for accelerated performance via differentiability

Innovation

• End-to-end differentiable molecular simulation framework
• Modular integration of ML and classical potentials
• Efficient neighborlists and constraint algorithms