[论文解读] 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 Science被引用 22

一句话总结

JAX MD 是一个完全用 Python 实现的端到端可微分子动力学框架，利用 JAX 的自动微分和 JIT 编译技术，在 CPU、GPU 或 TPU 上加速模拟。它支持快速原型设计、神经网络的无缝集成以及对整个模拟过程的梯度计算，显著降低了计算科学领域研究人员的门槛。

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.

研究动机与目标

• 通过在 Python 中实现端到端可微分模拟，解决传统分子动力学软件复杂且低层次编码的问题。
• 消除 MD 工作流中对手动推导导数和低层次 C++/CUDA 代码的需求。
• 利用 JAX 的自动微分和 XLA 编译技术，通过 JIT 编译在 CPU、GPU 或 TPU 上加速 MD 模拟。
• 使研究人员能够对整个模拟过程计算梯度，并将神经网络直接集成到模拟流程中。

提出的方法

• 使用 JAX 的自动微分在纯 Python 中实现所有模拟逻辑，以计算模拟输出相对于输入的梯度。
• 利用 JAX 的 JIT 编译在不更改源代码的情况下，加速 CPU、GPU 或 TPU 上的数值计算。
• 通过 JAX 的函数式编程范式将 MD 模拟表示为纯函数，从而实现自动微分和优化。
• 通过在 Python 中定义的可微分势能函数，支持成对和多体相互作用。
• 通过将神经网络视为模拟图中的可微分组件，实现其与 MD 模拟的无缝集成。
• 提供开源实现，并附带交互式 Colab 笔记本，以展示核心功能。

实验结果

研究问题

• RQ1 分子动力学模拟能否在完全使用 Python 实现的同时，通过硬件加速保持高性能？
• RQ2 自动微分在完整 MD 模拟中的应用效果如何，能否有效计算物理可观测量的梯度？
• RQ3 在可微分模拟框架下，神经网络在 MD 工作流中的集成程度如何？
• RQ4 与传统的手工优化 C++/CUDA 实现相比，纯 Python 的 JAX 基础 MD 框架性能表现如何？
• RQ5 研究人员是否能通过消除低层次代码的复杂性，实现更快的迭代和实验周期？

主要发现

• JAX MD 实现了完全在 Python 中的端到端可微分分子动力学模拟，在通过 JAX 的 JIT 编译后，性能可与手工优化的 C++ 和 CUDA 代码相媲美。
• 该框架允许用户对整个模拟过程计算梯度，从而实现对初始条件、力场参数或神经网络组件的优化。
• 用 Python 编写的模拟代码保持了人类可读性和易修改性，与低层次语言相比，显著降低了实验的门槛。
• 神经网络与 MD 模拟的集成是无缝的，支持可微分架构搜索以及模型与动力学的联合优化。
• JAX MD 支持在 CPU、GPU 和 TPU 上运行，仅需极少的代码修改，实现了可移植性和硬件加速，无需编写低层次代码。
• 开源发布配合交互式 Colab 笔记本，加速了框架的采用，并在多样化的研究领域中实现了可复现性。