[論文レビュー] Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians

Ishan Amin, Sanjeev Raja+1 more|ArXiv.org|Jan 15, 2025

Model Reduction and Neural Networks被引用数 5

ひとこと要約

本論文は、基盤モデルから高速で専門化されたMLFFを作成するためのヘシアンベースの知識蒸留を提案し、MDシミュレーションにおける精度を保ちつつ大幅な速度向上とエネルギー保存を実現する。

ABSTRACT

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.

研究の動機と目的

• Fast, specialized MLFFsが downstreamタスクで物理的に健全性を保つ必要性を動機づける。
• 基盤MLFFからエネルギーヘシアンを蒸留して、小さくて高速なMLFFへ転写するKDフレームワークを提案する。
• 複数の基盤モデル、データセット、および下流の化学空間にわたってアプローチを実証する。
• 専門化されたMLFFが教師モデルを上回るか、同等にする一方で大規模な推論速度アップを達成できることを示す。

提案手法

• 基盤モデルのエネルギーヘシアンを専門データサブセット上で事前計算する。
• 小さなstudent MLFFを訓練し、ジョイント損失（エネルギー/力の一致＋教師へのヘシアン整合）を最小化する。
• ヘシアン監視のコストを削減するためにヘシアン行をサブサンプリングで取得する。
• full Hessiansを作成せずに、ベクトル-ヤコビ行列積を用いてヘシアン行を効率的に抽出する。
• 直接力モデルを改善するために勾配ベースのエネルギー整合項をオプションで含める。
• 複数のデータセットとタスクにわたり、ヘシアンKDをbaselines（undistilled, n2n, a2a）と比較する。

Figure 1: Proposed Hessian distillation schematic. In our proposed distillation approach, we start with a machine learning force field (MLFF) foundation model (FM) that has been trained on a large quantity of diverse data. We precompute energy Hessians of the FM over a specialized data subset. We th

実験結果

リサーチクエスチョン

• RQ1 ヘシアンベースの蒸留は物理的な健全性を犠牲にすることなく高速・専門化されたMLFFを生成できるか？
• RQ2 ヘシアン KDはノード特徴蒸留（n2n）や他のベースラインと比較して精度とMD安定性にどう影響するか？
• RQ3 ヘシアン監視をサブサンプリングして訓練コストを削減しつつ性能を維持できるか？
• RQ4 蒸留モデルは専門化された下流タスクで元の基盤モデルを上回りつつ、 substantialな速度アップを提供できるか？

主な発見

• 基盤モデルから蒸留された専門化MLFFは元のFMに比べて推論を最大20x高速化する。
• 蒸留モデルは専門タスクでFMの性能と同等またはそれを上回ることが多く、蒸留なしのベースラインを上回ることがある。
• ヘシアン蒸留はエネルギー/力のMAE、MD安定性、エネルギー保存、および幾何最適化をベースラインと比較して改善する。
• ヘシアンのサブサンプリング（たとえs=1でも）で精度を保ちつつ訓練コストを大幅に削減する。
• より大きいJMP-L FMからの蒸留は、未蒸留FMよりNVE MDシミュレーションにおけるエネルギー保存を改善する。

Figure 2: Energy Conservation in NVE MD Simulations of Buckyball Catcher. We plot the change in the model predicted energy over the trajectory for 5 independent initial conditions. Some simulations become unstable before 100 ps (denoted by $\times$ ). (a) Hessian distillation improves the energy con