PFD-kit: 预训练原子模型微调和蒸馏工作流

背景

机器学习力场兼具效率和精度，可对大尺度材料体系进行精确的原子模拟，在高熵合金、固态电池、聚合物、表界面防腐等材料研发领域有重大应用潜力。然而，对大多数潜在用户来说，训练机器学习力场仍然充满了挑战，其本质困难在于需通过DFT等昂贵的第一性原理计算方法对大量材料构型进行计算，方可构建完备的训练集。以当下流行的DPGEN主动学习方法为例，构建完备训练集需进行大量DFT计算和多轮多模型迭代训练，消耗极大的算力、人力和时间成本，阻碍了机器学习力场的商业应用。更为严重的是，实际工业应用的材料体系往往涉及无序、缺陷、表界面等复杂的构型，训练成本和复杂度大大超出一般用户可承受范围。

预训练原子大模型的出现有望解决机器学习力场的应用瓶颈。预训练大模型参数量庞大，通过海量DFT计算数据训练，具有强大的泛化能力和精度潜力，保存了大量可在不同材料体系和预测任务间迁移的知识。因此，用户可从预训练模型参数出发，通过模型微调，继承预训练模型的可迁移知识，大幅减少训练下游模型所需DFT计算。DFT计算数量减少不仅可降低训练集机时消耗，更可简化训练流程，大幅降低包含算力、人力和时间资源在内的全方位成本。基于预训练模型的潜力，我们开发了PFD-kit工作流，为用户提供一个基于预训练原子模型，快速、可靠、全自动生成终端模型的解决方案。

技术思路

PFD-kit包括微调和蒸馏等两个主要步骤，可由预训练模型产生用于大尺度、高精度原子模拟的终端模型。PFD-kit的后端工作流基于dflow开发，工作流基本流程如下：

图1. PFD-kit的工作流基本流程.

用户提供初始构型文件，并指定预训练模型，输入相关参数，即可全自动运行并自动回收输出模型。初始构型文件和预训练模型文件，均可从用户的Bohrium云存储直接上传，无需本地提交。

模型微调

微调工作流中，首先对初始构型进行微扰，并基于微扰结构做DFT单点能计算或短AIMD模拟构建初始微调训练集和微调模型。随后，工作流将以迭代探索和训练方式向微调训练集中添加新样本，反复执行上述步骤直至达到收敛标准。此方法能在尽可能减少DFT计算的同时，保证微调训练集的完备性。

模型蒸馏

微调模型继承了预训练原子大模型的庞大参数量，并不适于对大尺度原子模拟，故需以微调模型出发进行模型蒸馏，训练可精确描述用户指定体系，但参数量大幅减少、运算效率大幅提高的简化力场。蒸馏过程中，将由微调模型通过MD模拟采并标注大量数据，理想情况下，通过单次训练即可得到具有理想精度的蒸馏模型。

APP界面及接口设计

PFD-kit APP采用Launching框架开发，专为复杂的科学计算工作流研发，包含模型微调与模型蒸馏两个子任务模块。以微调工作流界面为例，用户可直接从Bohrium的个人或项目存储空间中直接指定结构文件、预训练原子模型参数文件和DFT计算所欲赝势/固态基组文件等，无需从本地提交。