迈向快速、专用的机器学习力场:通过能量 Hessian 矩阵蒸馏基础模型
Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians
AI总结:
本文提出基于能量Hessian矩阵的知识蒸馏方法,将通用MLFF基础模型的知识迁移到更小、更快的专用模型上,在保持甚至提升精度的同时实现最高20倍加速,并确保测试时分子动力学模拟的能量守恒。
摘要:
基础模型(FM)范式正在改变机器学习力场(MLFFs),利用通用表示和可扩展训练来执行各种计算化学任务。尽管 MLFF FMs 已开始缩小与第一性原理方法的精度差距,但对更快推理速度的需求依然强烈。此外,虽然研究越来越聚焦于在化学空间中迁移的通用模型,但从业者通常一次只研究一小部分系统。这凸显了对特定下游应用相关、快速且专用的 MLFFs 的需求,这些模型在保持训练时可扩展性的同时,维持测试时的物理合理性。在本工作中,我们提出了一种将通用表示从 MLFF 基础模型迁移到专用于化学空间特定区域的更小、更快的 MLFFs 的方法。我们将该方法表述为知识蒸馏过程,其中较小的“学生”MLFF 被训练以匹配“教师”基础模型能量预测的 Hessian 矩阵。我们的专用 MLFFs 可比原始基础模型快 20 倍,同时保留甚至在某些情况下超越其以及未蒸馏模型的性能。我们还表明,将具有直接力参数化的教师模型蒸馏到使用保守力(即作为势能导数计算)训练的学生模型中,成功利用了大规模教师的表示以提高精度,同时在测试时分子动力学模拟中保持能量守恒。更广泛地说,我们的工作提出了一种 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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。