[논문 리뷰] JAX, M.D.: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python

Samuel S. Schoenholz, Ekin D. Cubuk|arXiv (Cornell University)|2019. 12. 09.

Machine Learning in Materials Science인용 수 22

한 줄 요약

JAX MD는 CPU, GPU 또는 TPU에서 시뮬레이션을 가속화하기 위해 JAX의 자동 미분과 JIT 컴파일 기능을 활용하여 완전히 파이썬으로 구현된 엔드 투 엔드로 미분 가능한 분자 동역학 프레임워크입니다. 이는 빠른 프로토타이핑, 신경망의 원활한 통합, 전체 시뮬레이션을 거쳐 기울기 계산을 가능하게 하여 계산 과학 분야의 연구자들이 접근하기 쉬운 환경을 제공합니다.

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.

연구 동기 및 목표

• 기존 분자 동역학 소프트웨어에서 요구하는 복잡성과 저수준 프로그래밍을 해결하기 위해 파이썬에서 엔드 투 엔드로 기울기 가능한 시뮬레이션을 가능하게 하는 것.
• MD 워크플로우에서 수작업으로 유도한 도함수와 저수준의 C++/CUDA 코드가 필요 없도록 하는 것.
• JAX의 자동 미분과 XLA 컴파일을 활용하여 CPU, GPU 또는 TPU에서 JIT 컴파일을 통해 분자 동역학 시뮬레이션을 가속화하는 것.
• 연구자들이 전체 시뮬레이션을 거쳐 기울기를 취하고, 신경망을 직접 시뮬레이션 파이프라인에 통합할 수 있도록 하는 것.

제안 방법

• 모든 시뮬레이션 로직을 파이썬으로 구현하고, JAX의 자동 미분을 활용하여 시뮬레이션 출력에 대한 입력에 대한 기울기를 계산하는 것.
• JAX의 JIT 컴파일을 활용하여 소스 코드를 변경하지 않고도 CPU, GPU 또는 TPU에서 수치 계산을 가속화하는 것.
• JAX의 기능형 프로그래밍 패러다임을 활용하여 분자 동역학 시뮬레이션을 순수 함수로 표현함으로써 자동 미분과 최적화를 가능하게 하는 것.
• 파이썬에서 정의된 기울기 가능한 잠재 에너지 함수를 통해 상호작용을 이원 및 다체 상호작용 모두 지원하는 것.
• 신경망을 시뮬레이션 그래프 내의 기울기 가능한 구성 요소로 간주함으로써 분자 동역학 시뮬레이션에 신경망을 원활하게 통합하는 것.
• 인터랙티브 Colab 노트북을 제공하는 오픈소스 구현을 통해 핵심 기능을 시연하는 것.

실험 결과

연구 질문

• RQ1 하드웨어 가속을 통해 고성능을 유지하면서도 전체적으로 파이썬에서 분자 동역학 시뮬레이션을 구현할 수 있는가?
• RQ2 자기 미분이 전체 분자 동역학 시뮬레이션에 얼마나 효과적으로 적용되어 물리적 관측량의 기울기를 계산할 수 있는가?
• RQ3 기울기 가능한 시뮬레이션 프레임워크를 통해 얼마나 잘 신경망이 MD 워크플로우에 통합될 수 있는가?
• RQ4 기존 수작업 최적화된 C++/CUDA 대비 순수한 파이썬 기반의 JAX 기반 MD 프레임워크의 성능는 어떠한가?
• RQ5 저수준 코드의 복잡성을 제거함으로써 연구자들이 분자 동역학에서 더 빠른 반복 및 실험 주기를 달성할 수 있는가?

주요 결과

• JAX MD는 JAX의 JIT 컴파일을 통해 수작업 최적화된 C++ 및 CUDA 코드와 경쟁 가능한 성능을 보이며, 파이썬에서 완전히 엔드 투 엔드로 기울기 가능한 분자 동역학 시뮬레이션을 가능하게 합니다.
• 이 프레임워크를 통해 사용자는 전체 시뮬레이션을 거쳐 기울기를 계산할 수 있으며, 초기 조건, 힘 매개변수 또는 신경망 구성 요소를 최적화할 수 있습니다.
• 파이썬으로 작성된 시뮬레이션 코드는 인간이 이해하기 쉬우며 쉽게 수정 가능하여, 저수준 언어에 비해 실험에 접근하기 위한 장벽을 크게 낮춥니다.
• 신경망을 분자 동역학 시뮬레이션에 통합하는 것은 원활하며, 기울기 가능한 아키텍처 탐색과 모델 및 동역학의 공동 최적화를 가능하게 합니다.
• JAX MD는 코드 변경을 최소화하여 CPU, GPU 및 TPU에서 실행 가능하며, 저수준 프로그래밍 없이도 이식성과 하드웨어 가속을 제공합니다.
• 인터랙티브 Colab 노트북과 함께 오픈소스로 배포되어 다양한 연구 분야에서의 채택과 재현 가능성을 가속화합니다.