分享自:

GPUMD:一个用于构建精确机器学习势函数并执行高效原子模拟的软件包

期刊:J. Chem. Phys.DOI:10.1063/5.0106617

本文向读者介绍了GPUMD软件包的最新进展,特别是其核心机器学习势函数框架——神经进化势(Neuroevolution Potential, NEP)的最新版本NEP3。该研究由来自全球多所知名高校和研究机构的学者共同完成,包括渤海大学、阿尔托大学、北京科技大学、哈尔滨工业大学(深圳)、南京大学、香港大学、厦门大学、查尔姆斯理工大学、斯坦福大学和拉夫堡大学等。相关成果以论文《GPUMD: A package for constructing accurate machine-learned potentials and performing highly efficient atomistic simulations》的形式,于2022年发表在《Journal of Chemical Physics》期刊上。

研究背景与目标

在计算材料科学和凝聚态物理领域,原子尺度模拟是理解物质结构、性质与动力学行为的关键工具。传统的量子力学计算方法(如密度泛函理论,DFT)虽然精度高,但其巨大的计算成本严重限制了所能模拟的体系尺寸和时间尺度。经验势函数(如Tersoff势、嵌入原子法势)计算效率高,但精度和可移植性往往不足,难以描述复杂多样的原子环境和化学键合。机器学习势函数(Machine-Learned Potentials, MLPs)应运而生,旨在结合两者的优点:通过机器学习方法从高精度的量子力学计算数据中拟合出势能面,从而能以接近经验势函数的计算成本,达到接近量子力学的精度。

尽管已有许多MLP框架被开发出来(如QUIP-GAP、SNAP、AMP、DeePMD-kit等),但其计算速度通常比典型的经验势函数慢两到三个数量级。这限制了MLP在大规模、长时间尺度模拟中的应用。因此,开发兼具高精度和高计算效率的MLP框架,使其能够在普通科研人员可获取的计算资源(如单块GPU)上模拟百万乃至千万原子规模的体系,具有重要的科学意义和应用价值。

本研究的目标正是展示GPUMD软件包及其内置的NEP框架在实现这一目标方面的最新进展。具体而言,作者旨在:1)介绍NEP3在原子环境描述符方面的两项重要改进,以提升精度;2)通过多个材料体系的测试,系统对比NEP与其他主流MLP在精度和计算效率上的表现;3)展示基于NEP潜在空间(latent space)的主动学习(Active Learning, AL)方案,以高效构建训练数据集;4)通过一系列碳材料的模拟实例,演示GPUMD/NEP在实际科研问题中的应用能力;5)介绍与GPUMD配套的Python工具包(GPyUMD、Calorine、PyNEP),以促进其集成到现代科研工作流中。

详细工作流程

本研究并非一项单一的实验性研究,而是一项涵盖方法发展、性能评估、工具集成和应用演示的系统性工作。其工作流程可以概括为以下几个关键环节:

  1. NEP3方法学发展:这是研究的核心理论部分。作者在先前NEP1和NEP2版本的基础上,提出了NEP3。主要改进有两方面:

    • 径向函数改进:在NEP2中,径向描述符分量中的径向函数是单一的切比雪夫基函数。在NEP3中,作者将其改进为多个切比雪夫基函数的线性组合。这使得径向函数具有更强的表达能力,能够更好地区分不同原子类型,从而提升多组分体系的拟合精度。相关的可训练参数数量也相应增加。
    • 角向描述符扩展:NEP1和NEP2仅包含三体角向描述符分量。NEP3借鉴了原子簇展开(Atomic Cluster Expansion, ACE)方法的思想,引入了特定的四体和五体角向描述符分量。这使得对原子局部环境的描述更加完备,有助于提高回归精度。作者详细推导了这些高阶描述符分量的显式表达式,并进行了预计算优化,以确保计算效率。
    • 实现细节:作者给出了描述符向量总长度的计算公式,以及NEP3中总可训练参数数量的公式。同时,他们重新梳理了基于“分力”(partial force)公式的力、维里(virial)和热流表达式,确保了其适用于NEP这类多体势,并为高效GPU实现奠定了基础。
  2. 高效GPU算法实现:这是实现高性能计算的关键。NEP方法被完全集成在GPUMD这个基于CUDA C++编写的开源分子动力学模拟包中。为了实现极致效率,作者精心设计了GPU算法:

    • 线程-原子一一对应:每个CUDA线程负责计算一个原子的所有相关量(描述符、site energy、受力、维里等)。
    • 核心内核精简:整个NEP能量、力和维里的计算仅由4个CUDA内核(kernel)完成:第一个计算描述符并应用神经网络;第二个计算与径向描述符相关的力和维里;第三个计算与角向描述符相关的“分力”;第四个汇总计算与角向描述符相关的力和维里。这种设计保证了高度的并行性和算术强度。
    • 混合精度计算:原子位置、能量、力等输入输出使用双精度,而内核内部的大部分计算采用单精度,在保证精度的同时大幅提升计算速度。
  3. 软件包与工作流设计:GPUMD软件包包含两个主要可执行文件:nep(用于训练NEP模型)和gpumd(用于使用训练好的势函数进行模拟)。训练NEP模型需要准备三个输入文件:包含训练集结构和目标能量、力、维里的train.in;用于监控泛化能力的test.in;以及定义模型超参数和训练控制的nep.in。训练采用可分离自然进化策略(SNES)优化损失函数。训练完成后生成的nep.txt势函数文件可直接用于gpumd进行各种原子模拟。

  4. 性能基准测试与对比:为了系统评估NEP3的性能,作者在多个标准数据集上进行了测试,并与当前主流MLP进行了对比。

    • NEP3 vs NEP2:在MgAlCu合金数据集上,使用相同的计算资源,NEP3在能量、力和维里的均方根误差(RMSE)上均优于NEP2,证明了改进的描述符的有效性。
    • 硅通用数据集:使用一个包含多种相态硅的数据集,对比了NEP与基于线性回归的ACE方法。结果显示,仅使用三体描述符的NEP模型即可达到约0.1 eV/Å的力RMSE,而ACE需要用到六体描述符才能达到相似精度,且NEP的计算成本更低。这表明神经网络回归相比线性回归,可以用更小的描述符维度达到相当的精度。
    • 偶氮苯分子(rmd17数据集):对比了NEP与ANI、GAP、SGDML和线性ACE等模型。在达到相近精度(力平均绝对误差,MAE)时,NEP的计算速度比ANI快一个数量级以上,比GAP快两个数量级以上,也比线性ACE快数倍。
    • 碳多相数据集:进行了最全面的对比,参与者包括GAP、DeePMD、MTP、REANN和NEP。在精度方面,NEP模型的能量、力、维里RMSE处于中上水平,与MTP和REANN接近,优于GAP和DeePMD。在计算速度和内存效率方面,NEP展现出巨大优势:在单块Tesla V100 GPU上,NEP的模拟速度达到每秒数千万原子步,比DeePMD和REANN快一个数量级以上,比使用72个CPU核心的MTP快两个数量级。同时,NEP可模拟的最大原子数也远高于其他GPU加速的MLP。这项测试清晰地表明,NEP显著降低了精度-成本的帕累托前沿。
  5. 基于潜在空间的主动学习方案:为了减少构建训练集所需的第一性原理计算量,作者提出了一种新的主动学习方案。其核心思想是利用已训练NEP模型的潜在空间(定义为隐藏层神经元状态与输出层连接权重的乘积所张成的空间)作为不确定性度量。通过对潜在空间进行主成分分析降维,可以可视化并量化不同结构在描述符空间中的分布。通过选择与现有训练集在潜在空间中距离较远的新结构进行第一性原理计算并加入训练,可以高效地扩充数据集。作者以碳数据集为例,仅用786个精心选择的结构(远少于原始的4080个)训练出的NEP模型,在测试集上的性能与用全数据集训练的模型非常接近,证明了该方案的有效性和NEP方法的数据效率。

  6. NEP模型应用实例演示:为了展示GPUMD/NEP解决实际科学问题的能力,作者使用针对碳体系训练的NEP模型(NEP 4.2 Å)进行了一系列模拟:

    • 静态计算:计算了金刚石的晶格常数和弹性常数,结果与DFT-LDA参考值吻合良好,特别是拟合了维里数据后,弹性常数预测更准确。
    • 动态模拟
      • 拉伸断裂:模拟了金刚石在单轴拉伸下的应力-应变关系,观察到了脆性断裂,断裂强度约为200 GPa。
      • 熔化-淬火-退火:演示了通过高温熔化、快速淬火和低温退火制备非晶碳的过程,分析了温度、sp³键比例、径向分布函数和角分布函数的演化。
      • 热学性质:计算了非晶碳的振动态密度,并由此计算了量子修正后的热容。
      • 热导率:采用均匀非平衡分子动力学(HNEMD)方法及谱分解技术,计算了非晶碳在300 K下的热导率,并进行了量子修正。
  7. Python接口开发:为了提升易用性和集成度,作者开发了三个Python工具包:

    • GPyUMD:用于辅助生成GPUMD的输入文件(如xyz.in, run.in)和处理输出文件,并提供了与ASE(Atomic Simulation Environment)的接口。
    • Calorine:提供了更深入的集成,包括一个ASE计算器,允许用户直接从Python调用NEP模型计算能量、力和应力;以及一个通过pybind11实现的C++接口,可以访问GPUMD的内部数据结构,甚至提供CPU-only的计算器。
    • PyNEP:专门用于辅助基于NEP的主动学习工作流,可以方便地计算描述符和潜在描述符,并实现如最远点采样等结构选择算法。

主要结果与结论

本研究的主要结果可归纳如下: 1. 方法学进步:成功发展了NEP3框架,通过改进径向函数和引入高阶角向描述符,在不牺牲效率的前提下,进一步提升了MLP的精度,特别是在多组分体系中的表现。 2. 卓越的性能表现:通过一系列基准测试证明,GPUMD中的NEP实现,在保持高于平均水平的精度的同时,其计算速度和内存效率远超其他主流MLP(如DeePMD、GAP、MTP、REANN等)。这使得在单块高端GPU(如A100)上模拟千万原子规模、纳秒时间尺度的体系成为可能。 3. 高效的数据准备方案:提出的基于NEP潜在空间的主动学习策略,能够有效减少构建高质量训练集所需的第一性原理计算量,降低了MLP的应用门槛。 4. 强大的应用演示:通过碳材料的一系列算例,全面展示了GPUMD/NEP在计算结构性质、力学性质、相变过程和非晶材料热学性质等方面的实用性和可靠性。 5. 完善的软件生态:提供了GPUMD核心包及其Python接口(GPyUMD, Calorine, PyNEP),形成了一个用户友好、易于集成到现代计算材料学研究工作流中的软件生态系统。

基于以上结果,本研究得出结论:GPUMD软件包及其NEP框架是一个极具前景的工具,它使得高精度、大规模原子模拟不再仅限于少数拥有海量计算资源的机构,而能够被更广泛的研究群体所使用。NEP方法在精度、效率和易用性之间取得了出色的平衡。

研究的亮点与价值

本研究的亮点在于: * 性能突破:将机器学习势函数的计算效率提升到了前所未有的高度,显著降低了大规模高精度原子模拟的计算资源门槛。 * 方法创新:在描述符设计中巧妙地结合了基函数线性组合和高阶相关项,并配以极其高效的GPU实现算法。 * 系统整合:不仅发展了势函数模型,还构建了包含训练、模拟、后处理及主动学习在内的完整、高效、用户友好的软件工作流。 * 严谨验证:通过多个标准数据集和与多种主流MLP的全面、定量对比,扎实地证明了其方法的优越性。 * 开源与生态建设:完全开源,并积极构建Python生态,促进了工具的普及和社区发展。

这项工作的科学价值在于推动了机器学习势函数向实用化、大规模应用的发展。其应用价值则体现在为材料科学、物理、化学等领域的科研人员提供了一个强大的工具,可以以前所未有的规模和精度探索材料的力学、热学、相变等复杂性质,加速新材料的发现和设计。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com