DPA-1: 共建覆盖元素周期表的预训练大模型

近年来，机器学习势函数正在彻底改变了分子模拟领域的研究范式，大量基于第一性原理计算产生的数据极大拓展了模型的应用范围。然而，面对一个新的复杂体系，我们基本上仍然需要生成大量新的数据从头训练模型。参考人工智能其他领域的发展，能否利用大量已经产生的数据、复用训练好的模型，是减少模型生产成本亟待解决的难题。

近日，深势科技（DP Technology）以及北京科学智能研究院（AI for Science Institute，Beijing）研究员张铎、毕航睿等人和合作者在arXiv上预发表了名为《DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation》的文章[1]，通过对元素类型更优的编码以及利用关键的注意力机制，极大提高了Deep Potential之前版本模型[2]的容量和迁移能力，获得了覆盖元素周期表大多常见元素的大型预训练模型。在不同数据集上的迁移学习结果表明，模型能大幅降低新场景对数据的依赖。

项目背景

机器学习势函数在材料科学、计算物理等领域应用广泛，并取得了较大成果。然而，由于现有模型迁移能力不足、缺乏通用的大模型，面对一个新的复杂体系，要获得可用的、较为完备的势函数模型，科学家们基本上仍然需要获取大量计算数据并从头开始训练模型。随着电子结构数据的积累，类比计算机视觉（CV）或者自然语言处理（NLP）等其他人工智能领域的发展，“预训练+少量数据微调”是解决这个难题比较自然的想法。

为了实现这一范式，我们亟需一种具有强迁移能力、能容纳元素周期表大多数元素的模型结构。

DPA-1模型是基于DP系列模型的一次全面升级，利用关键的门控注意力机制（Gated Attention Machanism），对原子间的相互作用实现了更为充分的建模，通过在现有数据上的训练，能够学习到更多隐藏的原子交互信息，极大提升了模型在包含不同构象、不同组分的数据集之间的迁移能力，从而也提升了在数据生成时的采样效率；并且模型通过对元素信息的编码，拓展了对元素的容量。开发者将模型在含有56种元素的较大数据集上进行了预训练，并将此预训练模型在各种下游任务上进行了迁移学习，实验表明，此预训练模型能大幅降低下游任务训练所需数据量及训练成本、提高模型预测精度，从而对分子模拟相关领域产生深远的影响。

DPA-1的训练和分子动力学模拟功能均已在DeepModeling开源社区DeePMD-kit项目（https://github.com/deepmodeling/deepmd-kit）开源。相关工作在深势科技科学计算云平台Bohrium上完成。

模型介绍

DPA-1的模型结构如下图所示，

DPA-1引入了对元素的编码，对不同元素共用同一套网络参数，从而提升了元素容量；模型结构中最重要的是红色的门控注意力模块，类似于传统NLP领域的注意力机制，在原子局部环境矩阵上进行了类比多体相互作用的信息交互，并以正则化后的相对坐标点乘作为角度信息，对获得的注意力权重进行重新加权，从而实现类似门控的机制。

DPA-1在推理方面延续了DP系列模型的高效率，可以进行大规模原子、元素体系的分子动力学模拟。

实验案例

迁移性测试

为了测试DPA-1模型结构带来的迁移能力提升，我们人为将不同训练集划分成了多个子集，每个子集之间的组分、构型有较大差异（以AlMgCu为例，single子集中仅包含单质数据；binary仅有二元数据，即Al-Mg，Al-Cu，Mg-Cu；而ternary则是剩余的三元数据），我们在其中一些子集上训练，在另一些子集上进行测试，来考验模型在极端条件下的迁移能力。如下表所示，在不同的实验设置下，我们测试了DPA-1和DeepPot-SE（DeePMD的光滑版本）在AlMgCu合金、固态电解质（SSE）和高熵合金（HEA）数据集上的表现。

可以看到，对比DeepPot-SE，在某些条件下DPA-1的测试精度甚至能实现一两个数量级的提升，这说明模型可以从现有数据中学习到隐含的原子间交互信息，也进一步证明了模型强大的迁移能力。

预训练模型及迁移学习方案

为了实现“预训练+少量数据微调”的流程，我们对DPA-1设计了迁移学习的方案：首先在较大规模数据上进行预训练，然后根据新数据集的统计结果修改最后一层的能量偏差，便可作为新场景下的训练起点。我们在不同场景下验证了此流程的可行性，通过不断往训练集中添加新样本的方式，来反映模型的采样效率，如下图所示，

我们首先在AlMgCu数据集的一元和二元子集上进行了预训练（第一列），并在三元数据上测试，可以看到，DPA-1在仅有少量三元数据的场景下，也达到了较高的精度，对比DeepPot-SE可以节省大约90%的三元数据；

进一步，我们用DPA-1在包含56种元素的大型数据集OC2M上进行了预训练（绿色线），并将其迁移到了毫不相关的HEA和AlCu数据集，与从头训练的DPA-1（黄色）和DeepPot-SE（蓝色）进行了对比（后两列），可以看到，模型结构和预训练策略均可观地减少了下游数据依赖。

元素参数可视化

为了进一步研究此覆盖元素周期表大多数元素的预训练模型的可解释性，我们将模型中学习到的元素编码进行了PCA降维并可视化，如下图所示：

所有的元素在隐空间中呈螺旋状分布，同周期元素沿着螺旋下降，同族元素则垂直螺旋方向分布，巧妙地对应了其在元素周期表中的位置，也很好地证明了模型的可解释性。

未来展望

DPA-1的提出为机器学习势能函数生产打开了新的范式，证明了“预训练+少量任务微调”流程的可行性，未来我们将继续致力于势能函数自动化生产、自动化测试，也会继续关注比如多任务训练、无监督学习、模型压缩、蒸馏等操作，方便用户一键生成下游任务所需的势能函数。此外，更大更全的数据库、下游任务与dflow工作流框架的结合也是未来极具发展性的方向。

如何参与

欢迎各位小伙伴关注并参与DeepModeling开源社区的相关工作，感兴趣势能函数预训练模型开发的小伙伴可以通过GitHub与我们取得联系。

GitHub地址：

DeepModeling仓库：
https://github.com/deepmodeling

DeePMD-kit 项目地址：
https://github.com/deepmodeling/deepmd-kit

结语

机器学习势函数的发展，正是从“小农作坊”模式走向“大数据、大模型”安卓模式的时候，DPA-1的提出表明了基于大模型实现“预训练+少量数据微调”流程的可行性，这也是后续一系列工作的起点。同时，基于DeepModeling开源社区的合作模式，未来必将源源不断地涌现更多的探索，我们共同前进，共筑分子模拟的新未来。

参考资料

[1] DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. 

https://arxiv.org/pdf/2208.08236.pdf

[2] End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems.

https://proceedings.neurips.cc/paper/2018/file/e2ad76f2326fbc6b56a45a56c59fafdb-Paper.pdf

- End -

（如需转载图文请与公众号后台联系）

-------------------------------

推荐阅读