[Paper Review] OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials

Peter Eastman, Raimondas Galvelis+23 more|arXiv (Cornell University)|Oct 4, 2023

Machine Learning in Materials Science20 citations

TL;DR

OpenMM 8 adds support for machine learning potentials by integrating arbitrary PyTorch models to compute forces and energy, with a higher-level interface and optimized CUDA kernels, demonstrated on CDK8 and GFP in water.

ABSTRACT

Machine learning plays an important and growing role in molecular simulation. The newest version of the OpenMM molecular dynamics toolkit introduces new features to support the use of machine learning potentials. Arbitrary PyTorch models can be added to a simulation and used to compute forces and energy. A higher-level interface allows users to easily model their molecules of interest with general purpose, pretrained potential functions. A collection of optimized CUDA kernels and custom PyTorch operations greatly improves the speed of simulations. We demonstrate these features on simulations of cyclin-dependent kinase 8 (CDK8) and the green fluorescent protein (GFP) chromophore in water. Taken together, these features make it practical to use machine learning to improve the accuracy of simulations at only a modest increase in cost.

Motivation & Objective

• Introduce machine learning potentials into the OpenMM MD toolkit.
• Allow arbitrary PyTorch models to compute forces and energy within simulations.
• Provide a high-level interface for modeling molecules with pretrained ML potentials.
• Enhance performance with optimized CUDA kernels and custom PyTorch operations.

Proposed method

• Integrate arbitrary PyTorch models into OpenMM simulations to compute forces and energy.
• Provide a higher-level interface for users to model molecules with pretrained potentials.
• Implement optimized CUDA kernels and custom PyTorch operations to accelerate simulations.
• Demonstrate usability and performance on real systems (CDK8 kinase and GFP chromophore in water).

Experimental results

Research questions

• RQ1 How can arbitrary PyTorch models be integrated into OpenMM to compute forces and energy during MD?
• RQ2 What is the impact of ML potentials on the accuracy of MD simulations within OpenMM?
• RQ3 What performance improvements are achievable using optimized CUDA kernels and custom PyTorch ops?
• RQ4 How practical is the use of ML-based potentials for modeling specific biomolecular systems like CDK8 and GFP in water?

Key findings

• ML potentials can be integrated into OpenMM to compute forces and energy via PyTorch models.
• A higher-level interface enables easy use of pretrained ML potentials for molecular modeling.
• Optimized CUDA kernels and custom PyTorch operations substantially speed up simulations.
• Demonstrations on CDK8 and GFP in water show practical cost/accuracy benefits of ML potentials.