[论文解读] Distillation of atomistic foundation models across architectures and chemical domains

John L. A. Gardner, Daniel F. Thomas du Toit+9 more|ArXiv.org|Jun 12, 2025

Machine Learning in Materials Science被引用 3

一句话总结

论文提出了一种与架构无关的蒸馏协议，通过合成数据将大型原子级基础模型的知识转移到更小更快的学生MLIPs，在多样化的化学领域实现了 substantial 速度提升（从超过10倍到超过100倍不等）。该方法通过对多种架构进行蒸馏并使用较小的微调集，实现了在中等硬件上的准确、可扩展的分子动力学模拟。

ABSTRACT

Machine-learned interatomic potentials have transformed computational research in the physical sciences. Recent atomistic `foundation' models have changed the field yet again: trained on many different chemical elements and domains, these potentials are widely applicable, but comparably slow and resource-intensive to run. Here we show how distillation via synthetic data can be used to cheaply transfer knowledge from atomistic foundation models to a range of different architectures, unlocking much smaller, more efficient potentials. We demonstrate speed-ups of $> 10 imes$ by distilling from one graph-network architecture into another, and $> 100 imes$ by leveraging the atomic cluster expansion framework. We showcase applicability across chemical and materials domains: from liquid water to hydrogen under extreme conditions; from porous silica and a hybrid halide perovskite solar-cell material to modelling organic reactions. Our work shows how distillation can support the routine and computationally efficient use of current and future atomistic foundation models in real-world scientific research.

研究动机与目标

• 展示一个通用的蒸馏协议，将知识从原子级基础模型（FM）转移到更小、更新的MLIP，在化学领域具备广泛适用性。
• 通过对多种MLIP架构进行蒸馏并利用合成数据标注，证明架构无关的适用性。
• 量化计算效率与精度权衡，包括内存使用和MD稳定性，在具代表性系统中进行评估。
• 验证蒸馏模型通过基于MD的诊断和基准测试保留关键物理性质。
• 强调在中等硬件条件下使原子级FM更易于获取的现实意义。

提出的方法

• 在具有量子力学标签的领域特定结构子集上微调现有原子级FM。
• 使用微调后的FM通过rattle-relax-repeat增广在不进行MD模拟的情况下生成大规模合成数据集。
• 在合成数据上训练小型、快速的学生MLIP架构，以逼近FM的预测与标签。
• 在DFT测试集上评估蒸馏模型，并在MD模拟中比较结构与热力学性质。
• 展示在架构（TensorNet、PaiNN、ACE）及ACE/EDDP家族内的加速与可扩展性。
• 展示架构无关的与ASE计算器的兼容性，以及augment-atoms以实现端到端工作流。

实验结果

研究问题

• RQ1 合成数据蒸馏是否能够在不同架构之间将高容量原子级FM的知识转移到较小、较快的学生模型？
• RQ2 在保持相对于DFT标签的准确性前提下，可以实现多少速度提升与内存效率？
• RQ3 蒸馏的MLIP是否能够在多种化学领域的MD中重现关键的结构和动力学性质？
• RQ4 对于反应性和高能量配置，蒸馏的实际限制与领域边界是什么？
• RQ5 蒸馏结果如何随架构、截断半径和微调数据量变化？

主要发现

• 当从图网络FM迁移到其他图网络架构时，蒸馏可实现>10x的加速；在利用ACE框架时可实现>100x的加速。
• 蒸馏模型（TensorNet、PaiNN、ACE）在DFT标签上的力MAE接近微调FM，且MD速度优势显著。
• 蒸馏模型在单GPU上实现稳定的MD，并可扩展到超出FM内存限制的更大系统。
• 在各领域（水、氢、二氧化硅、MAPI，以及溶剂中的有机反应）中，蒸馏模型能够重现关键结构与动力学特征，某些指标甚至优于教师模型。
• 消融研究表明合成数据的扩展性提升了FM到DFT的准确性，蒸馏模型在比FM更小的截断半径下工作也不会显著损失精度。
• 该方法需要的领域数据较少（<50个DFT标记结构）用于微调，且完全使用开源工具实现自动化。