Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians

要約

基礎モデル (FM) パラダイムは、汎用表現とスケーラブルなトレーニングを活用して、さまざまな計算化学タスクを実行することで、機械学習力場 (MLFF) を変革しています。
MLFF FM は第一原理法との精度の差を縮め始めていますが、推論速度の高速化に対する強いニーズがまだあります。
さらに、研究は化学空間全体に転移する汎用モデルにますます重点を置いていますが、専門家は通常、特定の時点でシステムの小さなサブセットのみを研究します。
これは、トレーニング時のスケーラビリティを維持しながらテスト時の物理的な健全性を維持する、特定のダウンストリーム アプリケーションに関連する高速で特殊な MLFF の必要性を強調しています。
この研究では、汎用表現を MLFF 基礎モデルから化学空間の特定の領域に特化したより小型で高速な MLFF に転送する方法を紹介します。
私たちは、知識蒸留手順としてアプローチを定式化します。そこでは、より小規模な「生徒」MLFF が、「教師」基礎モデルのエネルギー予測のヘッセ行列と一致するように訓練されます。
当社の特殊な MLFF は、元の基礎モデルよりも最大 20 $\times$ 高速になり、その一方で、そのパフォーマンスと未蒸留モデルのパフォーマンスを維持し、場合によってはそれを超えます。
また、直接的な力のパラメータ化を使用した教師モデルから、保守的な力 (つまり、位置エネルギーの導関数として計算) でトレーニングされた学生モデルへの抽出により、エネルギーを維持しながら、大規模な教師からの表現をうまく活用して精度を向上できることも示します。
テスト時の分子動力学シミュレーション中の保存。
より広範には、私たちの研究は、MLFF 開発の新しいパラダイムを示唆しています。このパラダイムでは、基礎モデルが、一般的な化学サブセット用の小型で特殊なシミュレーション「エンジン」とともにリリースされます。

要約(オリジナル)

The foundation model (FM) paradigm is transforming Machine Learning Force Fields (MLFFs), leveraging general-purpose representations and scalable training to perform a variety of computational chemistry tasks. Although MLFF FMs have begun to close the accuracy gap relative to first-principles methods, there is still a strong need for faster inference speed. Additionally, while research is increasingly focused on general-purpose models which transfer across chemical space, practitioners typically only study a small subset of systems at a given time. This underscores the need for fast, specialized MLFFs relevant to specific downstream applications, which preserve test-time physical soundness while maintaining train-time scalability. In this work, we introduce a method for transferring general-purpose representations from MLFF foundation models to smaller, faster MLFFs specialized to specific regions of chemical space. We formulate our approach as a knowledge distillation procedure, where the smaller ‘student’ MLFF is trained to match the Hessians of the energy predictions of the ‘teacher’ foundation model. Our specialized MLFFs can be up to 20 $\times$ faster than the original foundation model, while retaining, and in some cases exceeding, its performance and that of undistilled models. We also show that distilling from a teacher model with a direct force parameterization into a student model trained with conservative forces (i.e., computed as derivatives of the potential energy) successfully leverages the representations from the large-scale teacher for improved accuracy, while maintaining energy conservation during test-time molecular dynamics simulations. More broadly, our work suggests a new paradigm for MLFF development, in which foundation models are released along with smaller, specialized simulation ‘engines’ for common chemical subsets.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google