Learning Fast and Accurate Machine Learning Force Fields via Joint Atomic Energy and Energy Hessian Distillation

ICLR 2026 Conference Submission16773 Authors

Keywords: graph neural network, machine learning force field, knowledge distillation

Abstract: Atomistic foundation models, trained on extensive and diverse datasets, now achieve near $\textit{ab initio}$ accuracy across broad molecular and material systems while demonstrating strong transferability across chemical spaces. However, their large parameter counts result in high inference latency and large memory requirements, hindering long-time-scale molecular dynamics simulations and deployment on resource-constrained hardware. In practice, researchers in physical chemistry often focus on specific chemical subdomains, where compact specialized models with fewer parameters would be sufficient—provided they inherit appropriate inductive biases from large foundation models. This need motivates distillation techniques that compress foundation models into efficient specialized models while preserving accuracy. In this paper, we propose an architecture-agnostic distillation method: Joint Atomic Energy--Energy Hessian Distillation. This approach augments state-of-the-art Hessian supervision with atomic energy, which complements low-frequency components at minimal computational overhead ($<$0.5\%). Compared with the current state-of-the-art method, our method consistently improves energy MAE over Hessian-only distillation (averaging 48.3\% on SPICE and 6.1\% on MPtrj datasets) while achieving comparable force MAE (average improvement of 1.4\%). Ultimately, our approach reduces parameter counts by 78\%–98\%, enabling fast and deployment-friendly specialized models for targeted chemical subdomains.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 16773