[论文解读] DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation

Duo Zhang, Hangrui Bi+4 more|arXiv (Cornell University)|Aug 17, 2022

Machine Learning in Materials Science被引用 27

一句话总结

DPA-1 引入基于注意力的 Deep Potential 模型，该模型在大规模多元素数据上预训练，并针对下游 PES 任务进行微调，达到更强的迁移性和样本效率。它还揭示了可解释的元素类型嵌入，与周期律趋势对齐。

ABSTRACT

Machine learning assisted modeling of the inter-atomic potential energy surface (PES) is revolutionizing the field of molecular simulation. With the accumulation of high-quality electronic structure data, a model that can be pretrained on all available data and finetuned on downstream tasks with a small additional effort would bring the field to a new stage. Here we propose DPA-1, a Deep Potential model with a novel attention mechanism, which is highly effective for representing the conformation and chemical spaces of atomic systems and learning the PES. We tested DPA-1 on a number of systems and observed superior performance compared with existing benchmarks. When pretrained on large-scale datasets containing 56 elements, DPA-1 can be successfully applied to various downstream tasks with a great improvement of sample efficiency. Surprisingly, for different elements, the learned type embedding parameters form a $spiral$ in the latent space and have a natural correspondence with their positions on the periodic table, showing interesting interpretability of the pretrained DPA-1 model.

研究动机与目标

• 构建一个可扩展、准确的PES模型，能够利用跨大量元素的大规模预训练。
• 开发一个保持能量守恒、平滑性和并行可扩展性的基于注意力的Deep Potential（DPA-1）。
• 展示在AlMgCu、SSE和HEA系统上的预训练到下游的迁移效率。
• 展示学习到的元素类型嵌入的可解释性及其与周期表的关系。

提出的方法

• 引入带有元素类型信息的局部嵌入以及对邻近原子施加自注意力机制。
• 通过转换的相对坐标矩阵和类型嵌入构建局部描述符，随后通过基于注意力的聚合来构建可训练的描述符。
• 使用Adam训练并结合能量/力损失的联合损失，配合用于平衡能量项和力项的调度方案。
• 通过调整最终能量偏置并选择性冻结参数，特别是类型嵌入，来对预训练模型进行微调。
• 在OC2M（56个元素）上进行预训练，并在下游任务上评估包括AlMgCu、SSE和HEA数据集。

实验结果

研究问题

• RQ1 一个经过预训练的基于注意力的PES模型能否在不同组成（单组分、二元、三元）和不同数据集之间泛化？
• RQ2 与从零开始训练相比，预训练是否降低下游任务所需的数据和计算资源？
• RQ3 学习到的元素类型嵌入是否展现出与周期表相关的物理可解释结构？
• RQ4 该模型能否在大型尺度的模拟中保持能量守恒、平滑性和可扩展性？

主要发现

• DPA-1在AlMgCu、SSE和HEA的验证准确性方面优于DeepPot-SE，当从有限组成数据转移时。
• 使用OC2M进行预训练可实现强大的样本效率提升，显著减少三元数据需求（例如在AlMgCu中，少用90%的三元样本）。
• OC2M预训练的DPA-1在OC2M验证上实现具有竞争力的能量/力MAE（0.681 eV, 0.076 eV/Å），相较于其他最先进模型。
• OC2M预训练的DPA-1中的元素类型嵌入大致呈螺旋状，与周期表分组相关，表明可解释的物理意义。
• DPA-1能够扩展到包含大量元素的数据集，而训练成本并未呈现平方增长，这与某些基线模型不同。
• 预训练方法带来更好的样本效率，并且保持能量守恒、适用于分子动力学的平滑力场。