新增高效训练NEP模型，开源的机器学习力场PWMLFF 2024.5版本上线

近几年来，一种结合了物理学、高性能并行计算和机器学习算法的新的科研范式——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版本）。

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。此外，我们对于多体描述符采用了与两体描述符相同的截断半径。

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

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

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

开源仓库：

https://github.com/LonxunQuantum/PWMLFF_library/tree/main/NEP_test

百度云网盘:

https://pan.baidu.com/s/1beFMBU1IehmNEpIQ9B8ybg?pwd=pwmt

[2] DOI：10.1145/2001576.2001692