Samuel Schoenholz1
Google Brain1
The last decade has seen explosive progress in machine learning. This rapid growth has been powered by expressive automatic differentiation frameworks that can execute programs efficiently on accelerators (like GPUs and TPUs) and can be orchestrated from high-level languages like Python. Rewriting traditional tools, such as molecular dynamics, in these frameworks can substantially improve workflows in Materials Science, allow us to seamlessly integrate machine learning models into simulations, and open qualitatively new research directions. This talk will describe, and give an update on, JAX MD which is an end-to-end differentiable molecular dynamics simulation library. We will focus on discussing the creative and unique ways in which automatic differentiation has been combined with traditional physics simulation to enable significant improvements in our ability to model and design materials.