龙讯旷腾PWmat团队研发的开源机器学习力场PWMLFF升级，新增高效训练NEP模型

news2024/11/21 0:30:49

近几年来，一种结合了物理学、高性能并行计算和机器学习算法的新的科研范式——AI for science迅速崛起，并为解决精度与尺度无法并存的问题带来了曙光。基于机器学习力场（machine learning force field, MLFF）的分子动力学，可以在保持第一性原理计算精度的同时，极大地扩展模拟的空间和时间尺度。

PWMLFF是一款由龙讯旷腾团队开发的开源软件包，采用GNU许可协议，旨在快速生成与从头算分子动力学精度相媲美的机器学习力场。我们最新推出PWMLFF 2024.5版本，相比于PWMLFF 2024.3版本，新版本增加了神经演化势（Neuro evolution potential，NEP）模型，包括基于梯度的分层卡尔曼滤波( Recurrent layer extended kalman filter，LKF)和自适应矩估计（Adaptive moment estimation，ADAM）优化器训练、lammps分子动力学接口（GPU、CPU版本）。

1.为什么引入NEP模型

NEP模型是GPUMD软件包中的一种机器学习势函数（2022b NEP3 GPUMD )[1]，其描述符借用了集团展开的思想，最高支持展开到5体相互作用，具有非常高的精度；而且开发团队通过多项式展开把多体相互作用显示表达出来，极大提高了计算效率。GPUMD中训练NEP采用了可分离自然演化策略（separable natural evolution strategy，SNES)[2]，由于不依赖梯度信息，实现简单。但是对于标准的监督学习任务，特别是深度学习，更适合采用基于梯度的优化算法。我们在PWMLFF 2024.5版本实现了NEP模型（NEP4，网络结构如图1所示），能够使用PWMLFF中基于梯度的LKF或ADAM优化器做模型训练。我们在多种体系中比较了LKF和SNES两种优化方法的训练效率，测试结果表明，LKF优化器在对NEP模型的训练中展现了优越的训练精度和收敛速度。NEP模型的网络结构只有一个单隐藏层，具有非常快的推理速度，而引入LKF优化器则大幅提高了训练效率。用户可以在PWMLFF中以较低的训练代价获得优质的NEP并使用它进行高效的机器学习分子动力学模拟，这对于资源/预算有限的用户非常友好。

图 1. NEP网络结构，不同类型的元素具有独立但结构相同的子神经网络。此外，与文献中NEP4网络结构不同的是，对于每层的bias，所有的子网络不共享最后一层bias。此外，我们对于多体描述符采用了与两体描述符相同的截断半径。

2.测试1-优化器对比

我们对多种体系进行了测试，所有测试中将数据集的80%作为训练集，20%作为验证集。我们在公开的HfO2训练集（包含包含𝑃21/c、𝑃bca、𝑃ca21和𝑃42/nmc相的2200个结构）上对NEP模型分别在LKF和演化算法（SNES, GPUMD）训练，它们在验证集上的误差下降如下图2中所示。随着训练epoch增加，基于LKF的NEP模型相比于SNES，可以更快收敛到更低误差（误差越低精度越高）。在铝的体系下（包括3984个结构）也有相似结果（图3）。此外，我们在LiGePS体系以及五元合金体系中也有类似结果，更详细数据可点击文末的训练和测试数据链接查看。

图 2. HfO2体系（2200个结构）下，NEP模型在LKF和SNES优化器下的能量（左图）和力（右图）收敛情况。图中虚线为SNES算法训练能够达到的最低loss水平。

图 3. Al体系（3984个结构）下，NEP模型在LKF和SNES优化器下的能量（左图）和力（右图）收敛情况。图中虚线为SNES算法训练能够达到的最低loss水平。

3.测试2-模型精度对比

深度势能（deep potential, DP）模型是目前广泛使用的一种神经网络模型，PWMLFF中实现了Pytorch版本的DP模型，该DP模型也可以使用LKF优化器。我们在多个体系下，使用LKF优化器对NEP模型和DP（PWMLFF）模型训练做了对比，结果如下图4中所示。在Al、HfO2、LiGePS（包含1万个结构）、[Ru、Rh、Ir、Pd、Ni]五元合金体系（包含9486个结构）下，PWMLFF中的NEP模型比DP模型收敛都更快，精度也更高。特别的，对于五元合金，我们采用type embedding DP以减少元素种类对训练速度的影响（在之前的测试中，我们发现，对五种以上的元素的情况，在PWMLFF的DP训练中引入type embedding可以获得比普通DP更高的精度）。