[Paper Review] JAX, M.D.: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python

Samuel S. Schoenholz, Ekin D. Cubuk|arXiv (Cornell University)|Dec 9, 2019

Machine Learning in Materials Science22 citations

TL;DR

JAX MD is an end-to-end differentiable molecular dynamics framework implemented entirely in Python, leveraging JAX's automatic differentiation and JIT compilation to accelerate simulations on CPU, GPU, or TPU. It enables rapid prototyping, seamless integration of neural networks, and gradient computation through entire simulations, significantly lowering the barrier for researchers in computational science.

ABSTRACT

A large fraction of computational science involves simulating the dynamics of particles that interact via pairwise or many-body interactions. These simulations, called Molecular Dynamics (MD), span a vast range of subjects from physics and materials science to biochemistry and drug discovery. Most MD software involves significant use of handwritten derivatives and code reuse across C++, FORTRAN, and CUDA. This is reminiscent of the state of machine learning before automatic differentiation became popular. In this work we bring the substantial advances in software that have taken place in machine learning to MD with JAX, M.D. (JAX MD). JAX MD is an end-to-end differentiable MD package written entirely in Python that can be just-in-time compiled to CPU, GPU, or TPU. JAX MD allows researchers to iterate extremely quickly and lets researchers easily incorporate machine learning models into their workflows. Finally, since all of the simulation code is written in Python, researchers can have unprecedented flexibility in setting up experiments without having to edit any low-level C++ or CUDA code. In addition to making existing workloads easier, JAX MD allows researchers to take derivatives through whole-simulations as well as seamlessly incorporate neural networks into simulations. This paper explores the architecture of JAX MD and its capabilities through several vignettes. Code is available at www.github.com/google/jax-md. We also provide an interactive Colab notebook that goes through all of the experiments discussed in the paper.

Motivation & Objective

• To address the complexity and low-level coding required in traditional molecular dynamics software by enabling end-to-end differentiable simulations in Python.
• To eliminate the need for handwritten derivatives and low-level C++/CUDA code in MD workflows.
• To accelerate MD simulations through JIT compilation on CPU, GPU, or TPU using JAX's automatic differentiation and XLA compilation.
• To allow researchers to take gradients through entire simulations and integrate neural networks directly into simulation pipelines.

Proposed method

• Implementing all simulation logic in pure Python using JAX's automatic differentiation to compute gradients of simulation outputs with respect to inputs.
• Leveraging JAX's JIT compilation to accelerate numerical computations on CPU, GPU, or TPU without changing the source code.
• Using JAX's functional programming paradigm to express MD simulations as pure functions, enabling automatic differentiation and optimization.
• Supporting pairwise and many-body interactions through differentiable potential functions defined in Python.
• Enabling seamless integration of neural networks into MD simulations by treating them as differentiable components within the simulation graph.
• Providing an open-source implementation with interactive Colab notebooks to demonstrate core capabilities.

Experimental results

Research questions

• RQ1 Can molecular dynamics simulations be implemented entirely in Python while maintaining high performance through hardware acceleration?
• RQ2 How effectively can automatic differentiation be applied to entire MD simulations to compute gradients of physical observables?
• RQ3 To what extent can neural networks be integrated into MD workflows using a differentiable simulation framework?
• RQ4 How does the performance of a pure Python, JAX-based MD framework compare to traditional hand-optimized C++/CUDA implementations?
• RQ5 Can researchers achieve faster iteration and experimentation cycles in MD by eliminating low-level code complexity?

Key findings

• JAX MD enables end-to-end differentiable molecular dynamics simulations fully in Python, with performance competitive to hand-optimized C++ and CUDA code when compiled via JAX's JIT.
• The framework allows users to compute gradients through entire simulations, enabling optimization of initial conditions, force parameters, or neural network components.
• Simulation code written in Python remains human-readable and easily modifiable, drastically reducing the barrier to experimentation compared to low-level languages.
• The integration of neural networks into MD simulations is seamless, allowing for differentiable architecture search and joint optimization of models and dynamics.
• JAX MD supports execution on CPU, GPU, and TPU with minimal code changes, enabling portability and hardware acceleration without low-level coding.
• The open-source release with interactive Colab notebooks accelerates adoption and enables reproducibility across diverse research domains.