JAX MD: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python
Authors: Samuel S. Schoenholz
Ekin D. Cubuk
Venue:  (2020)
Keywords: Automatic Differentiation
Software Library
Physics Simulation
Differentiable Physics

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.

Under review as a conference paper at ICLR 2020
JAX, M.D.
END-TO-END DIFFERENTIABLE, HARDWARE ACCELERATED, MOLECULAR DYNAMICS IN PURE PYTHON
Anonymous authors
Paper under double-blind review
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 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 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 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
along with an interactive Colab notebook that goes through all of the experiments discussed the paper.
INTRODUCTION
Understanding complex many-body systems is challenge that underlies many of the hard problems in the physical sciences. A ubiquitous tool at our disposal in trying understand such systems is to posit interactions between the constituents and then simulate the resulting dynamics. interactions can be identified such that the simulation captures macroscopic behaviors observed in experiments, then the simulation can be studied to gain insight into the physical system. Since one has access to the full microscopic state at each step, is possible test hypotheses and make measurements that would otherwise be impossible. Such techniques, generally called molecular dynamics (MD), have been used to understand a wide range of systems including molecules, crystals, glasses, proteins, polymers, and whole biological cells.
Significant effort has gone into number of high quality MD packages such as LAMMPS (Plimpton, 1995), HOOMD-Blut (Anderson et al., 2008; Glaser et al., 2015), and OpenMM (Eastman et al., 2017). Traditional simulation environments are large and specialized codebases written C++ or FORTRAN, along with custom CUDA kernels for GPU acceleration. These packages include significant amounts of code duplication and hand written gradients. The state of affairs is reminiscent of Machine Learning (ML) before the popularization of Automatic Differentiation (AD). Researchers trying a new idea often have to spend significant effort computing derivatives and integrating them into these large and specialized codebases. Simultaneously, the amount of data produced from MD simulations has been rapidly increasing, in part due to ever increasing computational resources along with more efficient MD software. Furthermore, deep learning is becoming popular tool

---

Under review as a conference paper at ICLR 2020
both for making MD simulations more accurate and for analyzing data produced in the simulations. Unfortunately, the issues facing MD libraries are even more pronounced when combining MD with deep learning, which typically involves complicated derivatives that can take weeks derive and implement.
Here introduce JAX, M.D. (JAX MD) which is new MD package that leverages the substantial progress made in ML software to improve this state of affairs. JAX MD is end-to-end differentiable, written in pure python, and is fast since simulations are just-in-time compiled to CPU, GPU, or TPU using XLA. Moreover, JAX MD is based on JAX (Bradbury et al., 2018; Frostig et al., 2018) which has strong neural network ecosystem that can be used seamlessly with simulations. In addition to a strong neural network ecosystem and just-in- just-in-time compilation, JAX can automatically vectorize caleulations over one- or multiple-devices. This makes easy to simulate ensembles of systems in JAX MD. We will begin with a short introduction to simulations in JAX MD followed by brief description of JAX and some architectural choices underlying JAX MD, We then explore the features of JAX MD through several experiments:
Efficient generation of ensembles of systems.
Using neural networks to do machine learning of a potential.
Meta-optimization through a simulation to optimize physical parameters.
Cooling a glass.
While these examples are designed to be illustrative, they similar to problems faced in actual research. Moreover, all but the last example would be significantly more difficult using existing tools.
JAX MD has so far implemented simple pairwise potentials (Lennard-Jones, soft-sphere, Morse) and the embedded atom method (EAM) (Daw & Baskes, 1984). It can also work with the Atomic Simulation Environment (Larsen al., 2017) and other first-principles calculations that can be accessed from Python (e.g. Quantum Espresso (Giannozzi et al., 2009)). Due to its efficient spatial partitioning strategy, can simulate millions of particles on single GPU, On the ML side, JAX MD has access to all of the MI. developments in JAX, including state-of-the-an convolutional networks and graph networks.
RELATED WORK
Automatic differentiation has enjoyed a rich history in machine learning as well as the physical sciences (Baydin et al., 2018). Backprop has been the core algorithm in the recent explosion of ML research, enabled by packages like TensorFlow (Abadi et al., 2016), Torch (Collobert et al., 2002), and Theano (Bastien et al., 2012). While still generally focused on ML, more recent packages have made automatic differentiation more generally available as a computational tool (e.g. Autograd (Maclaurin et al., 2015), JAX (Bradbury al., 2018; Frostig al., 2018), PyTorch (Paszke et al., 2017), Chainer (Tokui et al., 2015), and Zygote (Innes et al., 2019)).
In the physical sciences, automatic differentiation has been applied to large variety of problems in structural optimization (Hoyer et al., 2019), quantum chemistry (Tamayo-Mendoza et al., 2018), fluid dynamics (Müller & Cusdin, 2005; Thomas al., 2010; Bischof et al,, 2007), computational finance (Capriotti, 2010). atmospheric modelling (Charpentier & Ghemires, 2000; Carmichael et al., 1997), optimal control (Walther, 2007), physics engines (de Avila Belbute-Peres et al., 2018), protein modelling (Ingraham et al., 2018; AlQuraishi, 2019), and quantum circuits (Schuld et al., 2019). For further related work at the intersection of ML. and MD, please see Appendix E. Despite significant work the topic, the number of general purpose simulation environments that are integrated with AD is scarce.
3 WARM-UP: SIMULATING A BUBBLE RAFT
We begin with a lightning introduction to MD simulations. As an example, we're going to imagine some bubbles floating water that they live two-dimensional interface between water and air. We describe N bubbles by positions, (Fisiss- Since the bubbles are confined to the water's s surface, the positions will be 2-dimensional vectors, 7 E TR.. For simplicity, we assume that the
2

---

Under review as a conference paper at ICLR 2020
bubbles all are the same size and let their diameter be 1 without a loss of generality. We now have to posit interactions between the simulated bubbles that qualitatively model how real bubbles behave. More accurate interactions will produce more accurate simulations, which will in turn capture more realistic phenomena. For the purposes of this example, we assume that bubbles interact with each other in pairs. We model pairs of bubbles by defining an energy function for the pair that depends only on the distance between them. We will choose an energy that is zero if the bubbles aren't 't touching and then increases gradually as they get pushed together. Specifically, if mi the distance between bubble and j, we use a pairwise energy function,
U(ri)) - { riy)
if Tf <
(1)
rf >
Once an energy has been defined we can compute the forces on a bubble, Fú, as the negative gradient that forces move bubbles to minimize their energy. From
of the energy. From their definition, Eq. (1) low energy configurations will be those where bubbles are touching as little as possible. However, bubbles are situated on water which is full of water molecules that are moving around. These water molecules bump into the bubbles and push them slightly. To model the interaction between the bubbles and the water we will assume that there are very small random forces from the water that push the bubbles. This is model called Brownian motion and it is described by first-order differential equation relating the velocity of bubbles to the forces on them along with random kicks coming from the water,
dF,(1) = + V (E)
dt
Here (1) are forces, E. ~ NO 1) is i.i.d. Gaussian distributed noise, and k specifies the temperature of the water. Incidentally, this model of objects in water dates back to Einstein (1905).
In Appendix A we show an example where we define a function, final_position s simulationn (rng_key), that takes a random number generator state and returns the final positions of the particles after simulating for some time. In this example, although only simulated a small number of bubbles we were able to emulate a much large bubble raft by using what are known as "periodic boundary conditions" (which boundary conditions bubbles can wrap around the side of the simulation to the other side, this is a ubiquitous technique for simulating the "bulk" properties of a system. In Appendix A we also show figures from real experiment compared with the results of the simulation which shows striking similarities despite the significant simplifying assumptions we made in defining our simulation.
used in the popular game, "Asteroids"). With periodic
4 ARCHITECTURE
We begin by briefly describing JAX before discussing the architectural choices we made in designing JAX MD. JAX is the successor to Autograd and shares key design features. As with Autograd, the main user-facing API of JAX is in one-to-one correspondence with Numpy (Van Der Walt et al., 201 I), the ubiquitous numerical computing library for Python. On top of this, JAX implements sophisticated "tracing" machinery that takes arbitrary python functions and builds an Abstract Syntax Tree (AST) for the function called a "Jaxpr". JAX includes number of transformations that can be applied to Jaxprs to produce new Jaxprs. Examples of such transformations are: automatic differentiation (grad), vectorization on a single device (vmap), parallelization across multiple devices (pmap), and just-in- time compilation (jit). To see an example of this see Appendix D, which shows that the grad function ingests a function and returns new, transformed function that computes its gradient. This is emblematic of JAX's functional design; all transformations take functions and return transformed functions. These function transformations are arbitrarily composable so one can write e.g. jit (vmap (grad (2) ) to just-in-time compile a function that computes per example gradients of a function f. As discussed above, JAX makes heavy use of the accelerated linear algebra library, XLA, which allows compiled functions to be run in a single call on CPU, GPU, or TPU. This effectively removes execution speed issues that generally face Python programs.
JAX MD adopts similarly functional style with immutable data and first-class functions. In further departure from most other MD software, JAX MD features no classes or object hierarchies and instead uses data driven approach that involves transforming arrays of data. JAX MD often uses named
3

---

Under review as a conference paper at ICLR 2020
tuples to organize collections of arrays. This functional and data-oriented approach complements JAX's style and makes it easy to apply the range of function transformations that JAX provides. JAX MD makes extensive use of automatic differentiation and automatic vectorization to conciscly express ideas (e.g. force as the negative gradient of energy) that are challenging in more conventional libraries. Since JAX MD leverages XLA to compile whole simulations single GPU calls, it can be entirely written in Python while still being extremely fast. Together this means that implementing simulations in JAX MD looks almost verbatim like textbook descriptions of the subject.
While our use of JAX and XLA provides a significantly more productive research environment, it does have several drawbacks. Most significantly, the primitives exposed by XLA are often at odds with computations that are commonplace in molecular dynamics. Notably, XLA requires that shapes be statically known at compile-time and is often challenging to use complex data-structures. This is makes several operations, for example spatial partitioning using cell lists, suboptimal. While JAX MD) fast enough for many research applications, since XLA has been optimized for machine learning workloads JAX MD is still slower than production quality MD packages using custom CUDA kernels. Indeed, we benchmarked JAX MD against HOOMD Blue Anderson et al. (2008) on a 4096 particle, Lennard-Jones system and observe 1 112 step for HOOMD Blue and 3200 / step for JAX MD which represents a slowdown of around 25x We expect this gap to shrink as XLA (and other ML language frameworks such as MLIR) continue to improve along with JAX MD. We now go through the major systems underlying JAX MD.
4.1 SPACES
In MD we simulate collection of N particles in either two- or three-dimensions. In the simplest case, these particles are defined by collection of position vectors, [Fi]isieN- Some simulations are performed with F. IR. where D 2, 3 is the spatial dimension of the system; this is implemented in JAX MD using the space. free () function. However, as discussed in Section 3, it is more common to use periodic boundary conditions in which case 7 E IRD with the association that ri = Lêz for basis vectors e₁ and some "box size" L. In this case the simulation space is homeomorphic to D-torus; this is implemented in JAX MD using the space. periodic (box__s sizze)ee function.
These boundary conditions can be implemented by defining two functions. First, a function di: RD x R + RDD that computes the displacement between two partieles. This function can in turn be used to define metric on the space d TR. D x RDD - by d(ā, b) = [ā b) Note, that in systems with periodic boundary conditions d' computes the displacement between a particle and the nearest "image" of the second particle. Second, shift function μ RD R - RD must be defined that moves a particle by a shift. Motivated by this, in JAX MD we implement the spaces outlined above by functions that return a "displacement" and "shift" functions. We show an example below.
import jax, numpy as np from jax_md mm import space r_ np.array (10.5, 0, 01) 2_ mp.array ( (10, 0.5, 01)
di dispaceement shift = space. pericdicri (1.) dR displacement (r_1, r_2, t=0.1)
4.2 ENERGY AND FORCES
As discussed above, MD simulations often proceed by defining potential energy function, U(F between particles. The degree to which U((r))) approximates reality has a significant influence of the fidelity of the simulation. For this reason, approximating potential energy functions has received significant attention from the ML community. JAX MD allows potential energy functions to be arbitrary "JAX traceable" functions, E N XDD + R. including arbitrary neural networks.
JAX MD provides a number of predefined, common, choices of energy functions including several pair potentials- Lennard-Jones soft-sphere, and Morse well as the Embedded Atom (Daw & Baskes, 1984) many-body potential and soft-spring bonds. Functions to compute the energy of system are constructed by providing displacement function for example, energy_fun
4

---

Under review as a conference paper at ICLR 2020
energy. speerte__peeee__eeeeeee pair (displacement). Forces easily be computed using helper function quantity. force (energy Fun) which is thin wrapper around grad. In addition to the pre-defined energy functions, it is easy to add new potential energy functions to JAX MD. In Section 5.2 we show how to add a neural network many-body potential called the Behler- Perrinello (Behler, 2011). In sec. B we describe some additional tools provided by JAX MI) to easily define custom energy functions.
In many applications, the scalar function u(r) has compact support such that u(r) 0 Of To for some cutoff, Te E R. We say that particles are not interacting if Tij > To- The pairwise function defined in Eq. (4) wasteful in this case since the number of computations scales as O(N²) even though the total number of interactions scales as O(V). To improve calculations in this case we provide the function ce 1_ fun amap. cel1_ _el1_1ist (fun, box_ pox_azee r_C, example e e_posi i ons) that takes a function and returns new function that uses spatial partitioning to provide a speed up.
4.3 DYNAMICS AND SIMULATIONS
Once energy function has been defined, there are number of simulations that can be run. JAX MD supports simple constant energy (NVE) simulation as well as Nose-Hoover (Martyna al., 1992), Langevin, and Brownian simulations at constant temperature (NVT). JAX MD also supports Fast Inertial Relaxation Engine (Fire) Descent (Bitzek et al., 2006) and Gradient Descent to minimize the energy of systems. All simulations in JAX MD follow pattern that is inspired by JAX's optimizers: for simplicity, we will use Brownian motion as an example in this section. Simulations pairs of two functions: initialization function, state init__ut fun (key, positions), that takes particle positions and returns an initial simulation state and an update function, state update_ fun (state), that takes state and applies a single update step to the state. To see an example of this in the case of Brownian motion, see code from the warm-up up in sec. A. Simulation functions can also feature time-dependent temperatures or spaces in which case time parameter can be passed to the update function, state update_fur fun (state, t-t).
5 FOUR VIGNETTES
5.1 VECTORIZED GENERATION OF ENSEMBLES
Increases in computing power are increasingly due to device parallelism rather compute speed. Indeed GPUs are designed to process significant amounts of data in parallel and TPUs move futher in this direction by offering high speed interconnects between chips. This parallelism is often used to simulate ever larger systems. However, there other interesting uses of parallelism that have received less attention. Many of these methods (e.g. replica exchange MCMC sampling (Swendsen & Wang, 1986) nudged elastic band (Henkelman al., 2000)) involve simulating ensemble of states simultaneously.
Thanks to JAX, ensembling can be done automatically in JAX MD. For small systems, the amount of necessary compute can be sub-linear in the number of replicas since it can otherwise be difficult to saturate the parallelism of accelerators. Here go through an example where we automatic en- sembling to quickly compute statistics of simulation. Suppose we have function s simulate late (key) that simulates single system given random key and returns its final positions using code similar to Section 3. As discussed in Section 4 JAX includes the function vmap that automatically vectorizes computations. Here to run an ensemble of simulations we simply define,
vectorized. simulation vmap (simulat ion)
Fig. (a) shows some example simulations of small, 32-particle systems that were performed in parallel on a single GPU. These simulations are too small to saturate the compute on a single GPU and Fig (b) shows that the time-per-sine per-simulttion decreases with the number of simulations being performed in parallel. This scaling continues until a batch size of about 100 when the GPU compute becomes saturated.
5

---

Under review as a conference paper at ICLR 2020
10a
b
Secs / Simulation
10
10
10-2
10"
10a Vectorization Size
102
Figure I:
On the left are 6 of the configurations produced by the vectorized simulation function. On the right is the time-per-simulation using the vmap functionality of JAX.
5.2 EASY MACHINE LEARNED POTENTIALS
Historically energy functions were often derived by hand based on coarse heuristics and scarce experimental results were used to fit parameters. More recently, energy functions with a larger number of fitting parameters (e.g. ReaxFF (Van Duin et al., 2001)) have become popular due to their ability to accurately describe certain systems. However, these methods traditionally involve significant expert knowledge and fail for systems that deviate too much from those that they were designed for. A natural progression of this trend is to use neural networks and large datasets to learn energy functions. There were a number of early efforts that received mixed success; however, not until 2007 when Behler and Parrinello (Behler & Parrinello, 2007) published their general purpose neural network architecture that learned energy functions emerged as viable alternative to traditional approaches.
Since then large amounts work has been done on this topic, however the substantial progress in machine learned potentials has not seen as much use as might be expected. At the root of this discrepancy are two points of friction at the intersection of ML and MD that prevent rapid prototyping and deployment of learned energies. First, simulation code and machine learning code are written in different languages. Second, due to the lack of automatic differentiation in molecular dynamics packages, including neural network potentials in physics simulations can require substantial work which often prohibits easy experimentation (see Eq. 3 below).
To address these issues, several projects developed adapters (Artrith Urban, 2016; Artrith et al., 2017; Lot et al., 2019; Onat et al., 2018) between common MI languages, like Torch and Tensorflow, and common MD languages like LAMMPS. However, these solutions require researchers to be working in exactly the regime serviced by the adapter. One of the consequences of this is that the atomistic features which get fed into the neural network need to be differentiated by hand within the set of features easily take weeks months of
MD package to compute forces. Trying out work to compute and implement these derivatives.
As example, how JAX MD gets around these issues easily. The Behler-Parrinello architecture describes the total energy of the system as a sum over individual contributions from per-atom neural networks, U(FF.)) Ej ((F.)) where f(0: x) is a fully-connected neural network with parameters 0 and ((F)) are hand-designed, many-body, features for particle j. While many choices of features exist, one simple set are given by the local pair correlation function, (p) = Σ, 8(1, - p). which measures the the density of particles a distance p from a central particle. The Pahrir-Parrinello architecture can be described and initialized in two lines of python.
will fit neural network to the bubble potential defined in Eq. (1) and
init_fun,
stax, seriall (
stax. Dense (no (no_hidden bidden__nits) stax. Relu, # hidden layer 1 stax. Dense hed__hidde___mm units), stax, # hidden layer 2
Iense (1)) readout params init_fun (key, (-1, number_of_featiue

---

Under review as a conference paper at ICLR 2020
stax is JAX's native neural network library. It also easy to define the Behler-Parrinell loss using vmap and the JAX MD function pair _arr_corr fun quantity.pair www.:ppainccorrrttiiinniiiiiieei (dissaaceeessss as shown below.
seantity.pair_coreelleo lambda params, positions: np. spm (E (params, g (possitios))) a Eq. 4.
(displacements
def per_example loss (params, pesiticns):
return (0 (params, positicns) energy_ (positions) **2
def loss (params, batch patchh_ositions):
return np mean (vmap (per_ (per_example_ in_axis= (None, 0) (params, batch__ossiionss)
per jr__example_loss defines the MSE loss on a single state (atomic configuration) and loss is the total loss over a minibatch of states. We see comparison between the learned energies and ground truth energies after training the above architecture for 20 seconds on 800 example states in Fig 5.2
(a).
3.00
25 b
Predicted energy
2.75
20
P(cose)
2.50
15
2.25
10
2.00
5
1.75
0
2.00
2.25
2.50
2.75
3.00
0.5
0.6
0.7
0.8
0.9
Correct energy
cosê
Figure 2: Left panel shows the agreement between predicted energies and the correct energies of the bubble rafts in the test set. Right panel shows the distribution of the inner product between the correct force and the predicted force.
We now compute forces with JAX MD and consider how this would be implemented in standard MD package. Recall that F; = - ((F.)),, where is the potential energy of the system defined in eq (5.2). Thus,
F F= Σ dE(o; G. ((F.)) (GF)(F.)
(3)
(GF.)
dF.
using the chain rule. Since dE/dC, is the gradient neural network it easy get in most neural network packages and feed into MD. However, traditionally dG,/dr; dr; is pain point and has to be coded up by hand. In JAX, MD, we get dG, / /rr; for free without any extra work using JAX's grad function as grad (lambda parans,
E (params, g(r)), argnums=1)
This energy function and force can now be used in any JAX MD simulation. In Fig. 2 (b) we see comparison between the predictions of this network after 20 seconds of training on states generated in JAX MD. Despite the small amounts compute involved, we see reasonable agreement for energies and forces between the machine-learned potential and the real potential.
5.3 OPTIMIZATION THROUGH DYNAMICS
So far we have demonstrated how JAX MD can make common workloads easier. However, combining molecular dynamics with automatic differentiation opens the door for qualitatively research.
7

---

Under review as a conference paper at ICLR 2020
(BEiaD)
0.5
0.6
0.7
0.8
0.9
1.0
0.6
0.5
0.4
0.4
0.5
0.6
0.7
0.8
0.9
D
Figure 3: Panel a shows the average energy and the standard deviation of the energy at D. Panel shows the derivative we calculate by differentiating through energy optimization by gradient descent as a function of D.
One such avenue involves differentiating through the simulation trajectory to optimize physical parameters. There have been several excellent applications far in e.g. protein folding (AlQuraishi, 2019; Ingraham et al., 2018), but until now this has involved significant amounts specialized code. This vein of research is also similar recent work in machine learning meta-optimization (Andrychowicz et al., 2016; Metz et al., 2018).
We revisit the bubble raft example above. In this case, will control the structure of the bubble raft by differentiating through the simulation. As we saw in Section 3, bubble rafts form a hexagonal structure when all of the bubbles have the same size. However, when the bubbles have different sizes the situation can change considerably. To experiment with these changes, we' re going to set simulation of bubble raft with bubbles two distinct sizes. To keep things simple, we' 'll half of the bubbles have diameter and half have diameter D. To control the conditions of the experiment, we will keep the total volume of the bubbles constant (see appendix C. Unlike the previous simulations, we will minimize the energy of the system using function s simulate (di ameter, key) that returns the energy of system given a diameter and a random key. Using vmap we can vectorize the simulation to compute the statistics for ensemble of states at different diameters in parallel on the GPU.
Some example states at different diameters along with the energy function of diameter be found in Fig We see that the hexagonal structure breaks down in the two-species case. Moreover, that there is "most disordered" point when D 0.8, which be energy point in Fig. 3(a). The study of such disordered systems is often referred to as the study of "Jammed (O'hern& et al., 2003)" solids. However, this was somewhat brute-force way investigate the role of size-disparity in the structure of bubble rafts. Could we have seen the same result more directly? Since each energy calculation is result of differentiable simulation, we can differentiate through the minimization with respect to D. This would allow us to find extrema of the minimized-ennte energy function of diameter using first-order optimization methods. This could be implemented in JAX MD dE__D__fu fun grad (s i mul ate). Of course the de dk_di__ _d fun function can be vectorized to aggregate statistics from an ensemble.
the highest
The gradient is plotted in Fig. 3 (b). We see that the gradient is positive and constant for D < 0.8 corresponding to the linear increase the average energy. Moreover, we see that the derivative crosses zero exactly at the maximum average energy. Finally, we observe that the gradient goes back to zero at D = This suggests that D = 0.8 is the point of maximum disorder, as we found by brute force above. It also shows that D 1 is the minimum energy configuration of the diameter. Although we hadn't hypothesized it, we realize this must be true since D < states are symmetric with D
keep the total packing fraction constant.
8

---

Under review as a conference paper at ICLR 2020
-26000
-30400
-27000
-30500-
-30600
-28000
-30700
-29000
-30800
-30000
-30900
-31000
33100
-32000
-31000
10
10
10a
10a
10
10a
10
10a
10
102
10
10
10
Figure 4: On the left the average energy plotted against time for glasses cooled at different rates. On the right is the final energy function of the rate at which the glass was cooled.
5.4 COOLING A GLASS
The above examples have mostly focused on small scale problems that were pedagogical; however, JAX MD is fast enough to do research at scale. To demonstrate this consider a problem of ongoing study: the cooling of a liquid to form a glass. When liquids are cooled quickly enough, instead of crystallizing to form a solid they fall out-of-equilibrimm into glassy state. The transition of liquid into a glass is one of the longest standing problems in condensed matter physics (Angell, 1995b). One particularly interesting phenomenon within glass transition research is that the final properties of the glass depend on how slowly the glass itself is cooled (Angell, 1995a; Debenedetti & Stillinger, 2001: Schoenholz al., 2017). We consider glass known the Kob-Andersen Lennard-Jones mixture (Kob & Andersen, 1995) which continues to be widely studied and run several experiments cooling the liquid into a glass different rates. We plot the results this experiment in fig. 4 and we see that the final energy of a cooled glass gets lower it gets cooled for longer.
CONCLUSION
We have described recent work general purpose differentiable molecular dynamics physics package, JAX MD and demonstrated several instances where simplifies the research process and enables new avenues work.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In /2th (USENIX) Symposium on Operating Systems Design and Implementation ({OSDI) 16), pp. 265-283, 2016.
Mohammed AlQuraishi. End-to-end differentiable learning of protein structure. Cell systems, 8(4):
292-301, 2019.
Joshua Anderson, Chris Lorenz, and Alex Travesset. General purpose molecular dynamics simulations fully implemented on graphics processing units. Journal of computational physics, 227(10):5342-533 5359, 2008.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. In Advances in neural information processing systems, pp. 3981-3989 2016.
C Austen Angell. Formation of glasses from liquids and biopolymers. Science, 267(5206):1224-1995 1995a.

---

Under review as a conference paper at ICLR 2020
CA Angell. The old problems of glass and the glass transition, and the many new twists. Proceedings of the National Academy of Sciences, 92(15):6675-6682 1995b.
Nongnuch Artrith and Jörg Behler. High-dimensional neural network potentials for metal surfaces: A prototype study for copper. Physical Review B, 85(4):045439, 2012.
Nongmuch Artrith and Alexander Urban. An implementation of artificial neural-network potentials for atomistic materials simulations: Performance for tio2. Computational Materials Science, 1 11: 135-150, 2016.
Nongnuch Artrith, Tobias Morawietz, and Jörg Behler. High-dimensional neural-network potentials for multicomponent systems: Applications to zinc oxide. Physical Review B, 83(15):153101, 2011.
Nongnuch Artrith, Alexander Urban, and Gerbrand Ceder. Efficient and accurate machine-learning interpolation of atomic energies in compositions with many species. Physical Review B, 96(1): 014112, 2017.
Nongnuch Artrith, Alexander Urban, and Gerbrand Ceder. Constructing first-principles phase diagrams of amorphous li si using machine- leaning-ssssiieeddd algorithm. The Journal of chemical physics, 148(24)2441111,25) 2018.
assisted sampling with evolutionary
Albert Bartók, James Kermode, Noam Bernstein, and Gábor Csányi. Machine learning a general- purpose interatomic potential for silicon. Physical Review X, 8(4):041048, 2018
Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, James Bergstra, Ian Goodfellow, Arnaud Bergeron, Nicolas Bouchard, David Warde Farley, and Yoshua Bengio. Theano: new features and speed 2012.
Warde-Farley, improvements. arXiv preprint
Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of machine learning research, 18 (153), 2018
Jörg Behler. Atom-centeree symmetry functions for constructing high-dimensional neural network potentials. The Journal of chemical physics, 134(7):074106, 2011.
Jörg Behler and Michele Parrinello. Generalized neural-netword representation of high-dimensional potential-energy surfaces. Physical review letters, 98( 14):44, 2007.
Christian Bischof, H Martin Bücker, Arno Rasch, Emil Slusanschi, and Bruno Lang. Automatic differentiation of the general-purpose computational fluid dynamics package fluent. Journal of fluids engineering, 129(5):652-658, 2007.
Erik Bitzek, Pekka Koskinen, Franz Gähler, Michael Moseler, and Peter Gumbsch. Structural relaxation made simple. Physical review letters, 97(17):170211, 2006.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, and Skye Wanderman-Milne, JAX: composable transformations of Python+NumPy programs, 2018. URL http: / //gggtithbb//ggo///g com/ honn/goolle e/ jax.
Luca Capriotti. Fast greeks by algorithmic differentiation. Available SSRN 1619626, 2010.
Richard Car and Mark Parrinello. Unified approach for molecular dynamics and density- theory. Physical review letters, 55(22):2471, 1985.
Gregory R Carmichael, Adrian Sandu, et al. Sensitivity analysis for atmospheric chemistry models via automatic differentiation. Atmospheric Enviromment, 31(3))475-493 1997.
Isabelle Charpentier and Mohammed Ghemires. Efficient adjoint derivatives: application to the metcorological model meso-nh. Optimization Methods and Software, 13(1):35-63, 2000.
Stewart J Clark, Matthew D Segall, Chris J Pickard, Phil J Hasnip, Matt IJ Probert, Keith Refson, and Mike C Payne. First principles methods using castep. Zeitschrift für Kristallographie-Crsssse Materials, 22055/)):66-577, 2005.
10

---

Under review as a conference paper at ICLR 2020
Ronan Collobert, Samy Bengio, and Johnny Mariéthoz. Torch: modular machine learning software library. Technical report, Idiap, 2002.
Ekin D Cubuk, Samuel Stern Schoenholz, Jennifer M Rieser, Brad Dean Malone, Joerg Rottler, Douglas J Durian, Efthimios Kaxiras, and Andrea Liu. Identifying structural flow defects in disordered solids using machine-leaming methods. Physical review letters, 1 118010):08000, 2015.
Ekin D Cubuk, Samuel S Schoenholz, Efthimios Kaxiras, and Andrea J Liu. Structural properties of defects in glassy liquids. The Journal of Physical Chemistry B, 120(26):6139-614 6146, 2016.
Ekin D Cubuk, Brad D Malone, Berk Onat, Amos Waterland, and Efthimios Kaxiras. Representations in neural network based empirical potentials. The Journal of chemical physics, 147(2):024104, 2017a.
Ekin Dogus Cubuk, RJS Ivancic, Samuel S Schoenholz, DJ Strickland, Anindita Basu, ZS Davidson, Julien Fontaine, Jyo Lyn Hor, Y-R Huang. Y Jiang, et al. Structure property relationships from universal signatures of plasticity in disordered solids. Science, 358(63666 2017b.
Murray Daw and Michael I Baskes. Embedded-atom method: Derivation and application to impurities, surfaces, and other defects in metals. Physical Review B, 29(12):6443, 1984.
Filipe de Avila Belbute-Peree Peres, Kevin Smith, Kelsey Allen, Josh Tenenbaum, and J Zico Kolter. End-to-end differentiable physics for learning and control. In Advances Neural Information Processing Systems, pp. 7178-7189, 2018.
Pablo G Debenedetti and Frank H Stillinger. Supercooled liquids and the glass transition. Nature, 4 410(6825):259, 2001.
Volker L Deringer, Noam Bernstein, Albert P Bartók, Matthew J Cliffe, Rachel N Kerber, Lauren E Marbella, Clare Grey, Stephen R Elliott, and Gábor Csányi Realistic atomistic structure of amorphous silicon from wewediine--eaaninddirrie molecular dynamics. The journal of physical chemistry letters, 9(1 2811)2877-2885 2018a.
Volker L Deringer, Miguel A Caro, Richard Jana, Anja Aarva, Stephen R Elliott, Tomi Laurila, Gábor Csányi, and Lars Pastewka. Computational surface chemistry of tetrahedral amorphous carbon by combining machine learning and density functional theory. Chemistry of Materials, 30(21):
1438-74455 2018b.
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 13(7)xe1056659, 2017.
Albert Einstein. On the motion of small particles suspended liquids at rest required by the molecular-kinee theory of heat. Annalen der physik, 17:549-560 -60, 1905.
J Enkovaara, Carsten Rostgaard, J Jørgen Mortensen, Jingzhe Chen, M Dulak, Lara Ferrighi, Jeppe Gavnholt, Christian Glinsvad, V Haikola, HA Hansen, et al. Electronic structure calculations with gpaw: a real- space implementation of the projector augmented method. Journal of Physies: Condensed Matter, 22325), 2010.
Felix A Faber, Luke Hutchison, Bing Huang, Justin Gilmer, Samuel S Schoenholz, George E Dahl, Oriol Vinyals, Steven Kearnes, Patrick F Riley, and O Anatole Von Lilienfeld. Prediction errors of molecular machine learning models lower than hybrid dft error. Journal of chemical theory and computation, 13(1 (((11):5555553) 5264, 2017.
Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing, 2018
Paolo Giannozzi, Stefano Baroni, Nicola Bonini, Matteo Calandra, Roberto Car, Carlo Cavazzoni, Davide Ceresoli, Guido L Chiarotti, Matteo Cococcioni, Ismaila Dabo, et al. Quantum espresso: a modular and open-source software project for quantum simulations of materials. Journal of physics: Condensed matler, 21539),395502, 2009.

---

Under review as a conference paper at ICLR 2020
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning- Volluee 70, pp. 1263-1272. JMLR. org. 2017.
Jens Glaser, Trung Dac Nguyen, Joshua A Anderson, Pak Lui, Filippo Spiga, Jaime A Millan, David C Morse, and Sharon Glotzer. Strong scaling of general- purpose molecular dynamics simulations on gpus. Computer Physics Communications, 192:97-107, 2015.
Jürgen Hafner. Ab-initio simulations of materials using vasp: Density-functions theory and beyond. Journal of computational chemistry, (911)22044-220 2078, 2008.
Paul Z Hanakata, Ekin Cubuk, David K Campbell, and Harold S Park. Accelerated search and design of stretchable graphene kirigami using machine learning. Physical review letters, 121()5): 255304, 2018.
Graeme Henkelman, Blas P Uberuaga, and Hannes Jónsson. A climbing image nudged elastic band method for finding saddle points and minimum energy paths. The Journal of chemical physics, 113(22):9901-9900 2000.
Stephan Hoyer, Jascha Sohl-Dickstein, and Sam Greydanus. Neural reparameterization improves structural optimization. arXiv preprint arXiv::909.0422: 2019.
John Ingraham, Adam Riesselman, Chris Sander, and Debora Marks. Learning protein structure with a differentiable simulator. 2018.
Mike Innes, Alan Edelman, Keno Fischer, Chris Rackauckus, Elliot Saba, Viral B Shah, and Will Tebbutt. Zygote: differentiable programming system to bridge machine learning and scientific computing. arXiv preprint arxiv: 2019.
Walter Kob and Hans C Andersen. Testing mode modelcoupling theory for a supercooled binary lennard- jones mixture i: The van hove correlation function. Physical Review E, (16)):4626, 1995.
Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E Castelli, Rune Christensen, Marcin Dulak, Jesper Friis, Michael N Groves, Bjørk Hammer, Cory Hargus, et al. The atomic simulation environment-a python library for working with atoms. Journal of Physics: Condensed Matter, 29(27)2773002. 2017.
Ruggero Lot, Franco Pellegrini, Yusuf Shaidu, and Emine Kucukhenli. Panna: Properties from artificial neural network architectures. arXiv preprint arXiv:1907.03000 2019.
Xiaoguang Ma, Zoey $ Davidson, Tim Still, Robert JS Ivancic, SS Schoenholz, AJ Liu, and AG Yodh. Heterogeneous activation, local structure, and softness in supercooled colloidal liquids. Physical review letters, 122(2):028001, 2019.
Dougal Maclaurin, David Duvenaud, and Ryan P Adams. Autograd: Effortless gradients in numpy. In ICML 2015 AutoML. Workshop, volume 238, 2015.
Glenn J Martyna, Michael L Klein, and Mark Tuckerman. Nosé-hoover chains: The canonical ensemble via continuous dynamics. The Journal of chemical physics, 97(4):2635-2643, 1992.
Luke Metz, Niru Maheswaranathan, Brian Cheung, and Jascha Sohl-Dickstein. Meta-learning update rules for unsupervised representation learning. arXiv preprint arXiv: 1804,0020, 2018.
J-D Müller and P Cusdin. On the performance of discrete adjoint cfd codes using automatic differentiation. International journal for numerical methods in fluids, 47(8-9):939-945 2005.
Berk Onat, Ekin D Cubuk, Brad D Malone, and Efthimios Kaxiras. Implanted neural network potentials: Application to li-si alloys. Physical Review B, 97()):094106, 2018.
Corey S O' O'hern, Leonardo Silbert, Andrea Liu, and Sidney R Nagel. Jamming at zero temperature and zero applied stress: The epitome of disorder. Physical Review E, 68(1):01 1306, 2003.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
12

---

Under review as a conference paper at ICLR 2020
Steve Plimpton. Fast parallel algorithms for short-range molecular dynamics. Journal of computa- tional physics, 117 1995.
Pankaj Rajak, Rajiv K Kalia, Aiichiro Nakano, and Priya Vashishta. Neural network analysis of dynamic fracture layered material. MRS Advances, 109-11117-1177, 2019a.
Pankaj Rajak, Aravind Krishnamoorthy, Aiichiro Nakano, Priya Vashishta, and Rajiv Kalia, Structural
phase transitions in mowse 2 monolayer: Molecular dynamics simulations and variational autoencoder analysis. Physical Review B, 100(1):014108, 2019b.
Samuel S Schoenholz. Combining machine learning and physics to understand glassy systems. In
Journal of Physies: Conference Series, volume 1036, pp. 012021. IOP Publishing, 2018.
Samuel S Schoenholz, Ekin D Cubuk, Daniel M Sussman, Efthimios Kaxiras, and Andrea J Liu. A structural approach to relaxation in glassy liquids. Nature Physics, 12(5))469, 2016.
Samuel S Schoenholz, Ekin D Cubuk, Efthimios Kaxiras, and Andrea Liu. Relationship between local structure and relaxation in out-of-equilibriun glassy systems. Proceedings of the National Academy of Sciences. 114(2):263-2)) 2017.
Maria Schuld, Ville Bergholm, Christian Gogolin, Josh Izaac, and Nathan Killoran. Evaluating analytic gradients on quantum hardware. Physical Review A, 99(3):032331, 2019.
Atsuto Seko, Akira Takahashi, and Isao Tanaka. First-principles interatomic potentials for ten elemental metals via compressed sensing. Physical Review B, 9255):054113, 2015.
Austin D Sendek, Ekin D Cubuk, Evan Antoniuk, Gowoon Cheon, Yi Cui, and Evan Reed. Machine learning-assisted discovery of solid li-ion conducting materials. Chemistry of Materials, 31(2):322-32,, 2018.
Tristan A Sharp. Spencer L Thomas, Ekin D Cubuk, Samuel S Schoenholz, David J Srolovitz. and Andrea J Liu. Machine learning determination of atomic dynamics at grain boundaries. Proceedings of the National Academy of Sciences, 11944)):19443-100 2018.
José Soler, Emilio Artacho, Julian Gale, Alberto Gareía, Javier Junquera, Pablo Ordejón, and Daniel Sánchez-Portal. The siesta method for ab initio order-n materials simulation. Journal of Physics: Condensed Matter, 14(11):2745, 2002.
Qiming Sun, Timothy C Berkelbach, Nick Blunt, George H Booth, Sheng Guo, Zhendong Li, Junzi Liu, James D McClain, Elvira R Sayfutyarova, Sandeep Sharma, al. Pysef: the python-based simulations of chemistry framework. Wiley Interdisciplinary Reviews: Computational Molecular Science, 8(1):e 13400 2018
Daniel M Sussman, Samuel S Schoenholz, Ekin D Cubuk, and Andrea J Liu. Discomnectigg structure and dynamics in glassy thin films. Proceedings of the National Academy of Sciences, 114(40): 10601-10605, 2017.
Robert H Swendsen and Jian-Sheng Wang. Replica monte carlo simulation of spin-glasses. Physical review letters, 57(21):2607, 1986.
Teresa Tamayo-Mendoza, Christoph Kreisbeck, Roland Lindh, and Alán Aspuru-Guzik. Automatic differentiation in quantum chemistry with applications to fully variational hartree-fock. ACS central science, 4(5):559-566, 2018.
Jeffrey P Thomas, Earl Dowell, and Kenneth C Hall. Using automatic differentiation create nonlinear reduced-order-modke aerodynamic solver. AIAA journal, 48(1):19-24 2010.
Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: next-generation open
source framework for deep learning. In Proceedings of workshop on machine learning systems (LearningSys) in the twenty-ninth annual conference neural information processing systems (NIPS), volume 5, pp. 1-6, 2015.
Stefan Van Der Walt, Chris Colbert, and Gael Varoquaux. The numpy array: structure for efficient numerical computation. Computing in Science & Engineering, 13(2):22, 2011.
13

---

Under review as a conference paper at ICLR 2020
Adri CT Van Duin, Siddharth Dasgupta, Francois Lorant, and William A Goddard. Reaxff: reactive force field for hydrocarbons. The Journal of Physical Chemistry A, 105(41):9396-95)) 9409, 2001.
Andrea Walther. Automatic differentiation of explicit runge-kutta methods for optimal control. Computational Optimization and Applications, 36(1):83-108 2007.
Qian Yang. Carlos A Sing-Long, and Evan Reed. Learning reduced kinetic monte carlo models of complex chemistry from molecular dynamics. Chemical science, 8(8):5781-7766, 2017.
Kun Yao, John Herr, David W Toth, Ryker Mckintyre, and John Parkhill. The tensormol-0.1 model chemistry: A neural network augmented with long-range physics. Chemical science, 98)) 2261-2269, 2018.
A BUBBLE RAFT EXAMPLE CODE
Listing 1: A simulation function that takes random key and returns final particle positions.
N = 32
dt 1c-1 temperature 0.1
simulation_steet nr.arange (1000)
key randem. randem.PRNGKey (0)
box size box displacement, shift space.pericdic (box size)
box_size size numperumerrrrree www.uuuuuueuiiiieeee (psstityypattiiiee count-N, number cumber_densityyitt
energy_fun energy. s
def simulation (key): pos_key, sim key random. split (key)
R random uni fommform (pos_kep (N, 2), maxval-box_size
init_fn, apply_fn samulateebrowwiaa
evergy_fun shift, dt, temperature)
state init_fn (sim_key,, R)
do step 1ambda state, (apply_fn (state, t-t), t)
state, lax. scan (dc_ step, state, simulation_steet steps)
return State.position
positions simulation (key)
B DEFINING CUSTOM POTENTIALS
Many popular potential energy functions either pairwise bonded in the sense that
Upair ((r)))
Σ u(ri))
or Ubond ((ri))
Σ u(Faabia)
(4)
ISKJEN
1515B
where Ti is the distance between particles and and {bia]isisB indexes bond be- tween particles ba and bi2 for total of B bonds. In this case, we offer the functions energy_fun smap.pair gecalar_energy__ur
displacement_fun) and energy smap.bond bond (scalar (scalar_energyy fun, displacement fun) that will convert scalar function, u(r), to the either pair-potential defined in Eq. (4), An example of this is shown below.
from jax_md md import. smap
14

---

Under review as a conference paper at ICLR 2020
a
Figure 5: An experiment of a bubble raft compared with the results of simulation.
scalar_cnergy_fun lambda *waargs: r **
metric_fun space metric (displacement_fun)
pair_energy_fon smap.pa pair Seccalar_neery__aee metric_fun) E_pair _oir_energy_fue fun (positicns)
bonds arr.arry(110 111)
bond energy smap bond (scalar_ energy fun, metric fun, bonds) bone bend energy fun (positicns)
The difference between these two functions amounts to the choice of whether to use dpair or dbood defined above. The two-line examples above and in Section 5.2 should be contrasted with the significant undertaking that would required to implement these features in traditional MD packages.
C OPTIMIZATION THROUGH DYNAMICS
If there N bubbles then the total volume of water filled by bubbles is,
N π ( +1)
Voublle
where the factor of 8 comes from the fact that system is split into two halves and using diameters not radii. Since the volume of our simulation is V L if we want to keep the "packing
fraction". Vbabbles/ / constant then we will have to scale the size of the box to be, Vunite
D EXAMPLE JAXPR AND ITS GRADIENT
Listing Python
Listing 3: Jaxpr
def (x) return ** print (2.0) 8.0
Lambda ; let b pow a 3.0 in [b]
f(x) 23
Listing Python
Listing S: Jaxpr
from Jax.api. import grad df_dx grad print (df dx (2.0)) 12.0
lambda ;
df
let b pow a 2.0 о = mul 3.0 b d safe 1.0
3x
dx
in [d]
15

---

Under review as a conference paper at ICLR 2020
E FURTHER RELATED WORK
General MD packages have been widely used to simulate molecules and solids, either using first- principles potentials (using software packages that derive the potential from quantum mechanics (Car & Parrinello, 1985) e.g. Quantum Espresso (Giannozzi et al., 2009), VASP (Hafner, 2008), SIESTA (Soler al., 2002). GPAW (Enkovaara al., 2010), CASTEP (Clark et al., 2005), PySCF (Sun et al, 2018)) or with empirical potentials (using approximate potentials that descrive specific atomic interactions e.g. LAMMPS (Plimpton, 1995), HOOMD-BIIe (Anderson et al., 2008; Glaser et al., 2015), and OpenMM (Eastman et al., 2017)). HOOMD-BIus in particular has been built with GPU acceleration in mind from the beginning, with the ability to script MD experiments using Python.
Coupled with the growing interest in deep learning, machine learning (ML) has become a popular tool for analyzing data that is produced by MD (Cubuk et al., 2015; Schoenholz et al., 2016; Cubuk et al., 2016; 2017b: Schoenholz, 2018; Rajak et al., 2019b;a; Sharp et al., 2018; Sussman et al., 2017; Hanakata et al., 2018; Sendek et al., 2018; Yang et al., 2017; Ma et al., 2019), as well as making MD simulations faster and more accurate (Behler & Parrinello, 2007; Behler, 2011; Artrith al., 2011; Artrith & Behler, 2012; Artrith et al., 2018; Deringer et al., 2018a; Bartók et al., 2018; Yao et al., 2018; 8; Seko et al., 2015; Deringer et al., 2018b; Cubuk et al., 2017a; Faber et al., 2017; Gilmer et al., 2017).
16


Reviews:
[{"review_id": "SklDy3jycr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe 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.\n\nFurthermore, 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. ", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I did not assess the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "Hylmq1Ik9H", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI 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.  \n\nAll of these points are fine for a package-release/tutorial paper, but for a conference paper, might hope to see these addressed:\nDescription 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?).\nLimitations of the library and drawbacks of any design decisions (something must be traded at some point?) are not explicitly mentioned.  \nDespite mentioning numerous existing MD libraries, no performance comparison is drawn against any other.  \nCould 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.\n\nI took a quick glance at the code on github; it is substantial but not huge, cleanly organized, and is well-documented.  ", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "Hkga6iGecB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n\nThis 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. \n\nThe Github link is provided for reproducible research and future development. It should be encouraged.\n\nI am sure whether this paper fit the ICLR or not, or how deep learning community can benefit from it.\n\nThe 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.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "N/A", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

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