[논문 리뷰] Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians

Ishan Amin, Sanjeev Raja외 1명|ArXiv.org|2025. 01. 15.

Model Reduction and Neural Networks인용 수 5

한 줄 요약

이 논문은 기초 모델로부터 빠르고 특화된 MLFF를 만들기 위해 Hessian 기반 지식 증류를 도입하여, 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 $ imes$ 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.

연구 동기 및 목표

• 하위 태스크에서 물리적으로 건전성을 보존하는 빠르고 특화된 MLFF의 필요성을 제시한다.
• 에너지 Hessian을 기초 MLFF에서 더 작고 빠른 MLFF로 증류하는 KD 프레임워크를 제안한다.
• 여러 기초 모델, 데이터셋 및 하위 화학 공간에 걸쳐 이 접근법을 시연한다.
• 특화된 MLFF가 큰 추론 속도 향상을 달성하면서 교사 모델을 능가하거나 일치할 수 있음을 보인다.

제안 방법

• 특화된 데이터 하위 집합에서 기초 모델의 에너지 Hessians를 사전 계산한다.
• 작은 학생 MLFF를 학습시켜 합동 손실을 최소화한다: 에너지/힘 매칭과 교사에 대한 Hessian 정렬을 포함한다.
• 하위 샘플링을 통해 샘플링된 Hessian 행을 사용하여 Hessian 감독의 비용을 줄인다.
• 전체 Hessian을 만들지 않고 벡터-야코비 행렬 곱을 활용하여 Hessian 행을 효율적으로 추출한다.
• 직접 힘 모델을 개선하기 위해 그라디언트 기반의 에너지 일관성 항을 선택적으로 포함한다.
• 여러 데이터셋과 작업에서 Hessian KD를 베이스라인(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 Hessian 기반 증류가 물리적 건전성을 희생하지 않으면서 빠르고 특화된 MLFF를 낳을 수 있는가?
• RQ2 정확도와 MD 안정성에서 Hessian KD가 노드 특징 증류(n2n) 및 다른 베이스라인과 어떻게 비교되는가?
• RQ3 Hessian 감독 샘플링이 학습 비용을 줄이면서도 성능을 유지하는가?
• RQ4 증류된 모델이 특화된 하위 작업에서 원래의 기초 모델보다 더 나은 성능을 발휘하면서도 상당한 속도 향상을 제공할 수 있는가?

주요 결과

• 특화된 MLFF가 기초 모델에서 증류되어 원래의 FM보다 최대 20배 빠른 추론 속도를 달성한다.
• 증류된 모델은 특화된 태스크에서 FM의 성능에 필히 일치하거나 능가하는 경우가 많으며, 비증류 베이스라인보다 우수할 수 있다.
• Hessian 증류는 에너지/힘 MAE, MD 안정성, 에너지 보존, 기하 최적화를 베이스라인과 비교하여 향상시킨다.
• Hessian 서브샘플링(심지어 s=1)으로도 정확도를 유지하면서 학습 비용을 크게 줄인다.
• 더 큰 JMP-L FM으로부터의 증류는 undistilled 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