JAX MD: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python

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 github.com/jaxmd/jax-md along with an interactive Colab notebook.

Reviews and author responses:

Thank you for your careful review of our work and useful suggestions!

> Description of the elements of the design of JAX which are useful here are presented, and appear distinct from other 
> AD libraries like Tensorflow or PyTorch, although the authors stop short of explicitly stating which functionality 
> would be more difficult/impossible to support with the possible alternatives (automatic vectorization of the 
> simulations seems like one?).

We agree with your assessment that automatic vectorization would probably be the largest pain point associated with implementing JAX MD in a different AD library. Indeed, automatic vectorization is deeply integrated with JAX MD since a number of quantities are defined per-particle pair and then vectorized across systems of particles. Having said this, we believe that a TensorFlow version is probably possible since TF2 now supports “vectorize”. Despite our name, we have contemplated building a TF backend in a similar manner to Pyro.

> Limitations of the library and drawbacks of any design decisions (something must be traded at some point?) are not explicitly mentioned.  

This is a great question and we have added a discussion of this point to the text in section 4. The main tradeoff that we experience is that the primitives that XLA exposes are sometimes at odds with the most efficient primitives for a molecular dynamics simulation. This is particularly important for spatial partitioning where more complicated data structures are often used that are challenging to implement using XLA. We do have an implementation of a cell-list but it is complicated by the fact that shapes must be static in XLA. Our cell-list implementation uses a sort which scales like O(Nlog^2N) on GPU while standard cell-list implementations scale like O(N) when coded directly in CUDA. We have contemplated writing some custom code to circumvent these issues.

> Despite mentioning numerous existing MD libraries, no performance comparison is drawn against any other.  

This is also a great question (and related to the above). JAX MD is certainly slower than production quality / custom CUDA MD systems. We benchmarked JAX MD against HOOMD Blue on a very standard physical system using a V100 GPU in each case. We found that JAX MD took ~3200 microseconds / step while HOOMD Blue took ~112 microseconds / step. Thus, JAX MD appears to be around 25x slower. We believe that this performance gap will narrow as we improve infrastructure on our end along with improvements to XLA (and possibly MLIR). Having said this, JAX MD is fast enough to do many kinds of research with and where appropriate, we believe that the improvements to research productivity are worth the reduction in performance. We have added a discussion of this point to the text.

> Could also show some demonstration of running an experiment which has complexity on par with state of the art
> research?  The bubble raft example is great for illustrative purposes, but it could be better to save some of that for a 
> tutorial and use space to exercise this library on a relevant problem and show performance there.

Great idea! We have added a short discussion to the text cooling a liquid to form a glass. This represents an experiment of similar complexity to research being published today. 


Thank you for taking the time to review our work.

We would like to discuss the applicability of this work to the ML community. While it is true that JAX MD will be of use to Physicists, we note that there has been significant research among ML practitioners that would be aided by JAX MD. In particular, we note the following (non-exhaustive) list of papers [1-7] that were published recently in Machine Learning Conferences. In each case, these papers leverage physical simulation; however, they were hindered since the simulations used were not integrated with deep learning libraries. This is precisely the gap that JAX MD hopes to fill. We would like to draw particular attention to [1], “Learning Protein Structure with a Differentiable Simulator” that could have been implemented out-of-the-box using JAX MD and received an oral at ICLR last year. 

Apart from this, we believe (though there has been less work in this direction so far) that physical systems are an ideal environment to explore meta-optimization since the inner-loop is much better understood than neural networks in deep learning.

[1] Learning Protein Structure with a Differentiable Simulator
Ingraham et al.; ICLR 2019

[2] Interaction Networks for Learning about Objects, Relations and Physics
Battaglia et al.; NeurIPS 2016

[3] Visual Interaction Networks: Learning a Physics Simulator from Video
Watters et al.; NeurIPS 2017

[4] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions
Schütt et al.; NeurIPS 2017

[5] Learning Invariant Representations of Molecules for Atomization Energy Prediction
Montavon et al.; NeurIPS 2012

[6] A Compositional Object-Based Approach to Learning Physical Dynamics
Chang et al.; ICLR 2017

[7] End-to-End Differentiable Physics for Learning and Control
Belbute-Peres et al.; NeurIPS 2018


Thank you for taking the time to review our work and we appreciate your advice about the writing. We will fix the sentence you noted and generally make the writing more formal.

We would like to discuss the applicability of this work to the ML community. While it is true that JAX MD will be of use to Physicists, we note that there has been significant research among ML practitioners that would be aided by JAX MD. In particular, we note the following (non-exhaustive) list of papers [1-7] that were published recently in Machine Learning Conferences. In each case, these papers leverage physical simulation; however, they were hindred since the simulations used were not integrated with deep learning libraries. This is precisely the gap that JAX MD hopes to fill. We would like to draw particular attention to [1], “Learning Protein Structure with a Differentiable Simulator” that could have been implemented out-of-the-box using JAX MD and received an oral at ICLR last year. 

Apart from this, we believe (though there has been less work in this direction so far) that physical systems are an ideal environment to explore meta-optimization since the inner-loop is much better understood than neural networks in deep learning.

[1] Learning Protein Structure with a Differentiable Simulator
Ingraham et al.; ICLR 2019

[2] Interaction Networks for Learning about Objects, Relations and Physics
Battaglia et al.; NeurIPS 2016

[3] Visual Interaction Networks: Learning a Physics Simulator from Video
Watters et al.; NeurIPS 2017

[4] SchNet: A continuous-filter convolutional neural network for modeling quantum interactions
Schütt et al.; NeurIPS 2017

[5] Learning Invariant Representations of Molecules for Atomization Energy Prediction
Montavon et al.; NeurIPS 2012

[6] A Compositional Object-Based Approach to Learning Physical Dynamics
Chang et al.; ICLR 2017

[7] End-to-End Differentiable Physics for Learning and Control
Belbute-Peres et al.; NeurIPS 2018


This paper announces a new software package for simulating molecular dynamics which includes close integration with a neural network / machine learning library--the first to do so.  Straightforward access to hardware acceleration (e.g. GPU) is provided for both the simulation and machine learning.

I lean toward accepting this submission.  If it were only about simulation molecular dynamics using hardware accelerators, I would question the appropriateness of the venue, but because it is explicitly intended to support training and usage of learned potential functions, it seems suitable.  Still might be better placed in a physics/chemistry venue, as where most of the references come from and likely where users would, too. The application area is no doubt an important research technique.  The paper is clearly written, with enough specific examples to contrast previous pain points in this line of work against its smoother interface.  

All of these points are fine for a package-release/tutorial paper, but for a conference paper, might hope to see these addressed:
Description of the elements of the design of JAX which are useful here are presented, and appear distinct from other AD libraries like Tensorflow or PyTorch, although the authors stop short of explicitly stating which functionality would be more difficult/impossible to support with the possible alternatives (automatic vectorization of the simulations seems like one?).
Limitations of the library and drawbacks of any design decisions (something must be traded at some point?) are not explicitly mentioned.  
Despite mentioning numerous existing MD libraries, no performance comparison is drawn against any other.  
Could also show some demonstration of running an experiment which has complexity on par with state of the art research?  The bubble raft example is great for illustrative purposes, but it could be better to save some of that for a tutorial and use space to exercise this library on a relevant problem and show performance there.

I took a quick glance at the code on github; it is substantial but not huge, cleanly organized, and is well-documented.  

The paper presents a python package, called JAX MD for simulating molecular dynamics (MD). JAX MD provides automatic derivations and allows to easily incorporate machine learning models in the MD workflow.

The paper is clearly written and seems technically correct. However, given that I am a specialist of neither package implementation nor physics, I can not really asses that all the details are correct/useful.

Furthermore, even if this work will surely be of great use for the physics community, I am not not sure that the contribution of this paper is sufficient for ICLR. 



This paper describes a general purpose differentiable molecular dynamics physics package, JAX MD. It shows several instances, where it simplifies the research process and enables new avenues of work. 

The Github link is provided for reproducible research and future development. It should be encouraged.

I am sure whether this paper fit the ICLR or not, or how deep learning community can benefit from it.

The writing does not feel academic enough sometime. For example,  "Please let us know if there are features that you would find interesting. We are always seeking contributions!" Please consider the rephrase it.

Meta-review: The paper is about a software library that allows for relatively easy simulation of molecular dynamics. The library is based on JAX and draws heavily from its benefits.

To be honest, this is a difficult paper to evaluate for everyone involved in this discussion. The reason for this is that it is an unconventional paper (software) whose target application centered around molecular dynamics. While the package seems to be useful for this purpose (and some ML-related purposes), the paper does not expose which of the benefits come from JAX and which ones the authors added in JAX MD. It looks like that most of the benefits are built-in benefits in JAX. Furthermore, I am missing a detailed analysis of computation speed (the authors do mention this in the discussion below and in a sentence in the paper, but this insufficient). Currently, it seems that the package is relatively slow compared to existing alternatives. 

Here are some recommendations:
1. It would be good if the authors focused more on ML-related problems in the paper, because this would also make sure that the package is not considered a specialized package that overfits to molecular dynamics.
2. Please work out the contribution/delta of JAX MD compared to JAX.
3. Provide a thorough analysis of the computation speed
4. Make a better case, why JAX MD should be the go-to method for practitioners.

Overall, I recommend rejection of this paper. A potential re-submission venue could be JMLR, which has an explicit software track.

Decision: reject