Colloquium with Senior Research Scientist at Google Brain Sam Schoenholz

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. In this talk I will describe substantial recent advances in software that has taken place in machine learning. I will then go on to discuss how we can leverage these advances to improve simulations in Physics with a focus on MD. To that end I will introduce JAX MD, 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 to 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.

Bio: Sam is a Senior Research Scientist at Google Brain working at the intersection between Machine Learning and Physics. His work focuses on better understanding neural networks using techniques from statistical physics as well as applying advances in Machine Learning to physical systems. Sam received his PhD from the University of Pennsylvania where he used machine learning to study disordered materials and glassy liquids.