[論文レビュー] On the design space between molecular mechanics and machine learning force fields

Yuanqing Wang, Kenichiro Takaba+15 more|arXiv (Cornell University)|Sep 3, 2024

Machine Learning in Materials ScienceMaterials Science被引用数 3

ひとこと要約

この論文は、分子力学（MM）と機械学習力場（MLFF）の間の設計空間を探求し、現在のMLFFが正確ではあるが、広範な利用にはあまりに遅いと主張している。本稿では、単純で微分可能な演算とE(3)-同変な帰納的バイアスを組み合わせた次世代のMLFFアーキテクチャを提案し、高い速度と正確性を両立させ、生物学的分子系の効率的で物理的に意味のあるシミュレーションを可能にする。

ABSTRACT

A force field as accurate as quantum mechanics (QM) and as fast as molecular mechanics (MM), with which one can simulate a biomolecular system efficiently enough and meaningfully enough to get quantitative insights, is among the most ardent dreams of biophysicists -- a dream, nevertheless, not to be fulfilled any time soon. Machine learning force fields (MLFFs) represent a meaningful endeavor towards this direction, where differentiable neural functions are parametrized to fit ab initio energies, and furthermore forces through automatic differentiation. We argue that, as of now, the utility of the MLFF models is no longer bottlenecked by accuracy but primarily by their speed (as well as stability and generalizability), as many recent variants, on limited chemical spaces, have long surpassed the chemical accuracy of $1$ kcal/mol -- the empirical threshold beyond which realistic chemical predictions are possible -- though still magnitudes slower than MM. Hoping to kindle explorations and designs of faster, albeit perhaps slightly less accurate MLFFs, in this review, we focus our attention on the design space (the speed-accuracy tradeoff) between MM and ML force fields. After a brief review of the building blocks of force fields of either kind, we discuss the desired properties and challenges now faced by the force field development community, survey the efforts to make MM force fields more accurate and ML force fields faster, envision what the next generation of MLFF might look like.

研究の動機と目的

• ML力場の導入における現在のボトルネックである計算速度を改善すること、その上で高い正確性を維持すること。
• 従来の分子力学（MM）と現代の機械学習力場（MLFF）の間のギャップを埋めるために、両者の設計空間を探ること。
• MLFF開発における主な課題、すなわち速度、安定性、一般化能力、データ効率性を特定すること。
• 複雑な生物学的分子系を定量的な予測能力をもってシミュレートできる、高速かつ正確な次世代MLFFの実現を図ること。
• MLFFのトレーニングにおける広範な化学空間カバレッジを実現するため、コミュニティ全体での高品質・多様なデータセットの構築を提唱すること。

提案手法

• 正確性と効率性の両立を図る関数形と設計原則を特定するため、既存のMMおよびMLFFアーキテクチャを調査する。
• E(3)-不変関数を普遍的に近似可能な、単純で微分可能な演算（例：ドット積）に基づく新しいMLFFアーキテクチャを提案する。
• 回転および並進不変性などの物理的帰納的バイアスをニューラルアーキテクチャに組み込み、滑らかさと安定性を確保する。
• 自動微分を用いて力の計算を実装し、一般化されたテンソル加速フレームワーク（例：PyTorch、JAX）内にモデルを実装する。
• スケーラブルで物理的知識を組み込んだトレーニング戦略（例：カリキュラム学習、適応バッチ処理）を活用し、データ効率性を向上させる。
• 大規模かつ高品質なデータセットを用いて事前学習されたフォーデーションモデルを提唱し、化学空間全体にわたる少サンプルまたはゼロショット一般化を可能にする。

Figure 1: Overview of the design space between molecular mechanics (MM) and machine learning (ML) force fields.

実験結果

リサーチクエスチョン

• RQ1 現代の力場において、速度と正確性の主な設計的トレードオフは何か。そして、それらを最適化するにはどうすればよいか？
• RQ2 機械学習力場を著しく高速化しつつ、1 kcal/mol未満の化学的正確性を維持することは可能か？
• RQ3 最小限の計算コストでE(3)-不変エネルギー関数を普遍的に近似可能な、アーキテクチャ的選択と帰納的バイアスとは何か？
• RQ4 コミュニティ主導の高品質なデータセットがあれば、広範な化学空間にわたる一般化を、大規模な再トレーニングなしに実現できるか？
• RQ5 ボルツマン分布の直接的生成モデル（例：ボルツマンジェネレータ、拡散モデル）は、分子シミュレーションにおいて明示的な力場の必要性を完全に排除できるのか？

主な発見

• 現在のMLFFは、限定された化学空間において1 kcal/molの化学的正確性を達成しているが、分子力学力場と比べて桁違いに遅い。
• MLFFの導入における主なボトルネックは正確性ではなく計算速度であり、特に大規模な生物学的分子系のシミュレーションにおいて顕著である。
• 単純で微分可能な演算（例：ドット積）に基づく次世代MLFFアーキテクチャは、E(3)-不変関数を普遍的に近似可能でありながら、高い速度を維持できる。
• モデルアーキテクチャにE(3)-同変な帰納的バイアスを組み込むことで、物理的整合性、滑らかさ、安定性を確保しつつ表現力も損なわない。
• 少サンプルまたはゼロショット一般化を実現するためには、コミュニティ全体での高品質・多様なデータセットの構築が不可欠である。
• ボルツマンジェネレーターや拡散モデルなどの新たな生成モデルは、将来的にはボルツマン分布を一度のステップで直接サンプリングすることで、明示的な力場の存在を不要にする可能性がある。

Figure 2: Between MM and QM energies and forces, there is little correlation. Scatter plots and kernel density estimate (KDE) of: (a): MM energy ( $U_{\mathtt{MM}}$ , mean-subtracted) plotted against QM energy ( $U_{\mathtt{QM}}$ , mean-subtracted); (b): MM force magnitude ( $|F|_{\mathtt{MM}}$ ) pl