[論文レビュー] 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 は、JAX の自動微分および JIT コンパイル機能を活用して、CPU、GPU、TPU でシミュレーションを高速化する、Python で完全に実装されたエンド・ツー・エンドで微分可能な分子動力学フレームワークです。研究者が迅速にプロトタイプを開発でき、ニューラルネットワークをシームレスに統合でき、シミュレーション全体を通じて勾配計算が可能になるため、計算科学分野の研究者にとっての障壁が著しく低下します。

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 コンパイルを活用して、CPU、GPU、TPU での JIT コンパイルにより MD シミュレーションを高速化すること。
• 研究者がシミュレーション全体を通じて勾配を取得でき、ニューラルネットワークをシミュレーションパイプラインに直接統合できること。

提案手法

• シミュレーション出力の入力に関する勾配を計算するために、JAX の自動微分を活用して、すべてのシミュレーション論理を純粋な Python で実装すること。
• JAX の JIT コンパイルを活用して、ソースコードを変更せずに CPU、GPU、TPU での数値計算を高速化すること。
• JAX の関数型プログラミングパラダイムを活用して、MD シミュレーションを純関数として表現し、自動微分および最適化を可能にすること。
• Python で定義された微分可能なポテンシャル関数を通じて、ペアワイズおよび多数体相互作用をサポートすること。
• シミュレーショングラフ内の微分可能なコンポONENTとしてニューラルネットワークを扱うことで、MD シミュレーションへのシームレスな統合を可能にすること。
• インタラクティブな Colab ノートブックを提供するオープンソース実装を提供することで、コア機能のデモを可能にすること。

実験結果

リサーチクエスチョン

• RQ1 ハードウェア加速により、パフォーマンスを維持したまま、分子動力学シミュレーションを完全に Python で実装できるか？
• RQ2 自動微分が、物理的観測量の勾配を計算するために、MD シミュレーション全体にどの程度効果的に適用できるか？
• RQ3 微分可能なシミュレーションフレームワークを用いて、MD ワークフローにニューラルネットワークをどの程度シームレスに統合できるか？
• RQ4 純粋な Python で構築された JAX ベースの MD フレームワークのパフォーマンスは、従来の手作業で最適化された C++/CUDA 実装と比べてどの程度か？
• RQ5 低レベルのコードの複雑さを排除することで、研究者が MD のイテレーションや実験サイクルをより迅速に進められるか？

主な発見

• JAX MD は、JAX の JIT コンパイルを経由することで、手作業で最適化された C++ や CUDA コードと比較しても競争力のあるパフォーマンスを発揮する、完全に Python で実装されたエンド・ツー・エンドで微分可能な分子動力学シミュレーションを可能にします。
• このフレームワークにより、ユーザーはシミュレーション全体を通じて勾配を計算でき、初期条件、力のパラメータ、またはニューラルネットワークコンponents の最適化が可能になります。
• Python で記述されたシミュレーションコードは人間が読みやすく、簡単に変更可能であり、低レベル言語と比較して実験の障壁が著しく低下します。
• ニューラルネットワークの MD シミュレーションへの統合はシームレスであり、アーキテクチャ探索の微分可能化およびモデルとダイナミクスの共同最適化が可能になります。
• JAX MD は、最小限のコード変更で CPU、GPU、TPU での実行をサポートしており、低レベルのコーディングを伴わずにポータビリティとハードウェア加速を実現します。
• インタラクティブな Colab ノートブックを併せたオープンソースリリースにより、多様な研究分野における採用促進と再現可能性が加速されています。