HPC+AI驱动的第一性原理科学智能计算平台

摘要

目的

智能科学; 基于第一性原理的计算; 分子动力学分析; 主动学习策略; 卡尔曼滤波方法; 模型压缩技术

引言

基于量子力学的核心理论

以机器学习为代表的人工智能方法应对传统科学计算所面临的维度障碍提供了新的解决方案。研究结果表明，在处理高维空间函数时，神经网络的复杂度与蒙特卡洛方法相当复杂度[2-3]其表现可被广泛应用于科学计算领域。这一发现不仅为神经网络在多个AI for Science领域取得了成功案例[4⇓⇓⇓⇓-9]也为解决科学问题展示了巨大潜力。在第一性原理计算方面自2007年以来逐渐受到重视并得到广泛应用特别是在利用超级计算机实现数十亿原子体系在极短时间内完成模拟方面取得了显著进展[4,10]这些进展推动着第一性原理方法朝着解决现实世界中复杂科学问题的方向稳步前进

然而，在现有可模拟的时间与空间尺度下，在材料、化学、生物等领域中直接解决复杂现实科学问题仍存在数量级上的巨大鸿沟。要缩小这一差距，则一方面依赖于领域专用超级计算机等超算体系结构的进一步发展完善；另一方面则需要在科学应用与智能计算方法上持续突破

在研究范式和工作流层面的革新深刻改变了传统科学计算格局

其中首要问题是科学数据的应用方式发生根本转变：如何高效生成以及合理利用科学数据成为倍受关注的核心议题

其次新的工作流模式对AI训练样本利用效率提出了更高要求

同时对AI模型的泛化与外推能力提出了更高要求

此外还对模型训练时间和快速迭代需求响应能力提出了更高要求

针对第一性原理计算领域的变革与挑战

1 相关工作

Behler[11]最先提出了一种基于神经网络的方法来模拟高维系统的原子势能面函数，并将其理论基础建立在第一性原理精度的基础之上。该方法通过将总势能在各原子单独贡献的部分进行分解，并结合近邻效应模型来进行简化的计算处理。其创新之处在于能够实现任意尺寸复杂系统的高维势场建模能力，并为其在实际应用中提供了可靠的基础框架。伴随着人工智能及其深度学习技术的巨大发展，在这一领域取得了显著进步

2017年, Schütt等人开发了基于连续滤波卷积的SchNet[12]模型, 该模型通过引入光滑滤波函数使卷积层得以非网格化处理连续原子位置变化, 同时实现了平移与旋转不变性, 并生成势能面使其符合能量守恒定律. 该方法成功拟合出力场. 2018年, 张林峰等人提出了具备端到端设计并保持物理对称性的原子间势能面模型[13]. 该模型通过自动生成特征实现了平移、旋转及交换不变性, 达到了高效的端到端拟合效果. 近年来, DeePMD-kit[6]软件在超级计算机上实现了亿亿个原子体系每天一纳秒级的速度模拟[4], 大大推动了一阶原理分子动力学研究的重大跨越.

在分子建模领域中，图神经网络（Graph Neural Networks, GNN）因其与分子结构天然契合的优势而备受关注，并在第一性原理精度应用方面取得了显著进展。DimeNet++[14]通过采用消息传递机制结合原子对间距信息以及原子三联体角度信息来预测分子量子力学性质，并展现出卓越的效果；GraphNVP[15]则将变分自编码器（Variational Autoencoder, VAE）与图神经网络相结合以学习化学分子分布特征；CGCNN[16]则专注于晶体材料建模与预测任务中的应用研究。

近年来，在自然语言处理及图像处理领域的 Transformer[17] 基础上发展起来的大模型技术取得了重大的进展。此外，在第一性原理精度下的分子动力学模拟方面也获得了部分应用并取得了一定的发展。2021年，微软亚洲研究院团队提出了 Graphormer[18] 模型。传统基于图神经网络（GNN）的分子建模方法在深层网络结构上存在过度平滑的问题，这限制了网络规模以及其表达能力。然而 Graphormer 模型通过将图结构信息编码至 Transformer 的自注意力机制中得以有效克服这一挑战。随后于 2022 年，在深度势能模型的基础上引入元素类型编码后构建而成的大规模势能面预测器 DPA-1[7] 基于 Transformer 实现了类比多体相互作用下的原子间信息交互机制设计，并最终实现了对多种元素元素类型的预训练大模型构建以及在多个下游应用任务上的迁移学习能力培养。该研究成功地将"大数据预训练 + 少量特定任务数据微调"这一策略引入到了高精度的第一性原理分子动力学领域之中。

基于国家"材料基因工程"发展战略与大数据驱动的科研创新趋势的推动，在国内科研机构与高校纷纷建立了相关领域的数据库以支持第一性原理领域的科学智能计算工作。其中中国科学院物理研究所刘淼等博士建立的Atomly[19]材料科学数据库收录了超过3万种无机化合物的数据；北京大学潘峰教授团队开发了高精度材料科学数据库系统。此外由北京科学智能研究院(AISI)等单位共建的科学智能广场 (AIS-Square)[20]则是一个集科学智能数据共享、模型构建及工作流协同于一体的开源平台不仅涵盖了金属半导体氧化物以及过渡金属化合物等多种材料计算数据还配备了50余种专用科学智能模型及其配套的工作流系统。中国科学院计算机网络信息中心自主研发的科学数据银行（Science Data Bank）[21]作为覆盖各领域科学研究数据的一站式存储平台拥有超过34.28TB的数据量不仅显著提升了科研成果的价值也为科学研究与人才培养做出了重要贡献

在获取科学数据时存在较高的成本及有限的覆盖范围，在应用层面面临诸多挑战。因此，在实际应用中能够实现对目标系统的构型空间进行有效探索及精确标注的关键技术即为主动学习方法的应用场景之一，并逐渐成为科学数据分析的重要工具之一

针对科学智能模型的训练过程而言

2 HPC+AI 驱动的第一性原理科学智能计算平台概述

该科学智能计算平台基于HPC和AI技术架构设计，并通过图1展示了其整体框架。该平台采用了模块化层次结构设计，在软件与模型层、智能计算层以及科学数据层之间实现了有机整合。由三个关键环节构成：软件与模型层、智能计算层以及科学数据层。这些关键组件通过协同作用形成了完整的科学智能计算工作流。

图 1

图****1HPC+AI****驱动的第一性原理科学智能计算平台总体结构

Fig.1The overall architecture of a HPC-based intelligent systems integrated with foundational principles for scientific computation.

2.1 软件与模型层

软件与模型层主要由第一性原理计算软件和科学智能计算模型两部分构成。这些软件主要用于构建量子力学精度的领域模拟数据库库。这些数据库不仅可以作为训练科学智能模型的数据来源，还可以用来验证和评估这些模型的表现。平台采用的主要第一性原理计算软件包括PWmat[37-38]和DGDFT[39]等国产先进计算工具；其中使用的科学智能计算模型包括DeePMD-kit[6]、DPA-1[7]以及MLFF[9]等基于第一性原理构建的高精度力场模型。

2.2 智能计算层

智能计算层作为核心驱动力，在推动科学智能计算工作流中发挥着关键作用。其主要流程如下：

1. 利用第一性原理计算软件生成领域数据库，并对具体应用场景中的初始训练样本进行标注
2. 借助主动学习模块对目标场景的构型空间进行探索，并完成训练数据的采样与标注
3. 根据不同应用场景的需求，本系统采用高效优化器结合大规模并行训练方法，在较短时间内实现科学智能模型的快速构建与持续优化
4. 对训练完成后得到的科学智能模型进行压缩等推理优化处理，并将其部署在超级计算机上实现高效的推理运算

2.3 科学数据层

科学数据作为推动科学智能计算工作流的关键基础资源。在科学数据层中，核心构建是基于第一性原理的智能数据平台，其主要职责是承担全流程的数据管理和服务保障。在数据库组织方面，在第一性原理模拟过程中产生的原始物理属性和结构特征等关键信息被整合到第一性原理模拟数据库中；针对科学研究中常用的各类化学物质及其相关属性，在训练阶段生成的高质量样本被分类存储于科学智能模型训练样本采样数据库；专门用于存储优化后的各层次人工智能模型及其性能指标的系统则构成科学智能计算模型库；最后，在预测分析环节中被收集整理的各种目标参数及预测结果则归集于独立的科学智能预测数据库中。在工作流运行过程中，各阶段间的数据会持续进行双向交互：一方面科学研究通过调用预建的数据资源来推进工作进程；另一方面，在完成特定任务后生成的新类型、新特性的关键数据会被及时导入相关数据库积累沉淀，在不断迭代优化的过程中提升整体系统效能。

3 关键技术

3.1 第一性原理模拟数据生成与数据集制备

第一性原理精度分子动力学模拟的一个显著优势在于其数据标注过程完全依赖高性能超级计算机上的方程求解。 在处理不同物理化学问题时，所采用的基函数类型、是否引入伪势以及数值求解器的选择会有所差异。 例如，在量子化学领域中广泛使用的Gaussian软件基于高斯基函数展开；而VASP和Quantum Espresso等程序则常用于材料科学领域的模拟。 近年来我国学者在第一性原理计算软件开发方面取得了显著进展，并已推出了多个优秀的DFT与量子化学计算程序包如BDF[40]、PWDFT[41-42]等。 提名的是由于国产超算硬件架构的特点要求相关软件必须国产化才能发挥其性能优势 在本文工作中我们选择了PWmat[37-38]该程序并实现了对其国产曙光平台的支持

3.2 科学智能模型的高效训练

3.2.1 高效优化器

提升优化器是提高AI模型训练速度最直接的办法。

线性卡尔曼滤波模型[43]于1960年被提出，在此基础上通过将非线性系统线性化产生了非线性卡尔曼滤波EKF[44]，并进一步发展出了用于训练神经网络的GEKF[45]方法。相较于Adam，GEKF可以在保持精度的同时减少收敛所需的epoch数量，例如RuNNer软件包[46]利用GEKF方法训练三层全连接神经网络来预测H2O和Cu2S体系的能量和力，可以达到第一性原理精度。但是GEKF在降低收敛所需epoch数的同时，其在单个epoch中的计算量却较大，因此相较于Adam，其训练的总耗时反而高一个数量级。

为了降低计算量，在EKF的基础上又发展出了点解耦非线性卡尔曼滤波NDEKF[47]与层解耦非线性卡尔曼滤波LDEKF[48]。然而NDEKF和LDEKF等优化器在精度与求解速度方面做了折中，受限于精度要求，它们很难用于分子动力学模拟问题的训练。为此我们提出了RLEKF[9]优化器，该优化器在LDEKF的解耦合基础上对同层参数进一步解耦：若单层参数超过预先设定的解耦层大小，我们进一步均匀解耦该层，解耦层大小为预先设定值；若单层参数少于预先设定值则不做任何处理。相较于LDEKF，RLEKF采用了更均匀的解耦合策略，降低了卡尔曼滤波算法的计算复杂度，实现了更高的计算速度与精度，在科学智能计算模型的加速训练中取得了良好效果。

3.2.2 大规模并行训练

此外

3.3 基于卡尔曼滤波的主动学习策略

主动学习的数据采样过程通常通过图2来展示。图中各标记点代表了机器学习力场在MD模拟过程中遇到的不同构型状态。其中绿色标记的构型表明预测结果的不确定性处于可接受范围内；而蓝色标记则表明预测结果超出了可接受范围，则需要被选中进行标注并用于重新训练模型。这些红点则表明模型对相应构型的预测存在重大偏差，则需要被丢弃以避免影响后续建模效果。持续迭代这一过程直至模型能够覆盖足够的构型空间。其中用于评估预测结果不确定性的查询算法主要包括两类：一类是委员会查询[22]；另一类是单模型类型的查询方法（即单模型类型的查询方法），这些方法通过显式评估预测结果中的不确定性来改进模型性能。

图 2

图2主动学习的体系构型采集过程

Fig.2The system configuration acquisition process of active learning****

正如第3.2节所述，在模型训练过程中展示了极高的计算效率和精确度。基于RLEKF优化器并结合其中的协方差矩阵P提出了一种单模型不确定性度量算法KPU（Kalman Prediction Uncertainty）。该算法能够有效生成涵盖广泛构型空间具有代表性的训练数据集。KPU的具体定义如下：

使用RLEKF训练神经网络得到的权重近似分布：

Θ∼N(W,P)Θ∼N(W,P)

(1)

对f(Θ;xi)f(Θ;xi)做一阶线性近似，令：

Hi=Df(Θ;xi|w)Hi=Df(Θ;xi|w)和y_i=HiΘyi=HiΘ，且y_iyi服从正态分布：

y_{i∼N(HiW,HiPHTi)y}i∼N(HiW,HiPHiT)

(2)

则得到预测量y_i=f((Θ;xi)yi=f((Θ;xi)的近似分布：

y_{i∼N(f(W;xi),HiPHTi)y}i∼N(f(W;xi),HiPHiT)

(3)

由此定义模型对于预测的不确定度量：

KPU(xi)=HiPHTiKPU(xi)=HiPHiT

(4)

其中 i 为原子下标。

主动学习的过程由训练、探索与标注三个环节构成，并按照一定的循环迭代模式展开。当系统满足收敛条件时即停止运行：此时系统已充分地探索了构型空间，并准确地标记出一组具有代表性的数据样本；随后利用这些数据样本训练出一个精确的PES模型以供后续使用。为了保证这一流程的有效执行需具备相应的计算资源支持：其中用于模型训练任务而言需依赖高性能服务器或集群；而进行分子动力学模拟则需要大量地依赖于CPU资源；同时标记过程也需要配备充足的GPU资源以加速运算效率。值得注意的是在探索阶段以及不确定性分析等过程中各计算任务之间并无相互依赖关系因此可通过将任务分配至多个计算节点并行处理从而显著提升采样效率。为了实现对不同计算任务的有效调度并充分利用各类计算资源提升整体采样效率我们开发了一套基于KPU架构的主动学习平台ALKPU系统如图3所示该平台能够自动规划并执行迭代过程合理分配各计算任务所需的人力物力并将运算节点分散至多台独立设备上协同工作最后汇总分析结果输出优化建议。

图 3

图3主动学习平台ALKPU的训练、探索和标注模块数据流

Fig.3The data flow within the ALKPU active learning platform's training, exploration, and labeling modules.

3.4 科学智能模型的大规模高效推理

在DeepMDFD-kit模型框架中, 嵌入网络被设计为构建保证物理对称性特征的原子环境描述符, 其核心模块的具体架构可参考图4.c. 该网络通过将一组由中心原子与邻居间距离所定义的标量特征进行编码, 生成反映局部环境信息的空间滤波器矩阵_G_i_. 这一计算流程可直观地通过图4.e来理解.

图 4

图****4DeePMD-kit****模型的计算流程与主要数据操作流程

Fig.4计算流程及其核心数据处理步骤在DeepMOLD-kit模型中

性能分析结果显示，在模型运行过程中嵌入网络所需的计算资源消耗占总量的比例均超过95%，而嵌入网络推理环节所占运行时长的比例则达到90%以上。针对DeePMD-kit中的嵌入网络体系我们开发了一种基于分段五阶多项式拟合的模型压缩方法同时结合了核运算优化技术和冗余参数消除策略以显著提升了该软件包的整体性能效率。

3.4.1 分段五阶多项式模型压缩

我们开发了一种基于打表法（tabulation）的模型压缩算法，并将其应用于一维嵌入网络的设计与实现过程中；该算法的核心思想在于，在完成训练的一维嵌入网络中，在其输入取值范围内的不同子区域内分别采用分段五阶多项式来进行模型拟合；在推理过程中，则需要将待处理变量对应于每个输出维度，并利用各自对应的区间内多项式系数计算出相应的多项式拟合值以获得最终输出结果。

通过模型压缩技术显著减少了该类模型的总浮点运算量（减少幅度超过80%），同时保证了推断误差维持在介于10^{-10}至10^{-11}之间，并成功达到了基于严格量子力学原理的分子动力学模拟所需精度水平。

3.4.2 Kernel 融合优化

在基准模型的基础上应用模型压缩算法后，在图4(e)中所示的计算_G_i_以及_R_i_T_G_i_矩阵乘的操作流程发生改变：通过查找表获得_G_i_矩阵的具体行信息后直接与_R_i_T_矩阵对应列的数据进行外积运算操作以获得_R_i_T_G_i_-相关分量的结果这一过程在并行机制下实现即各个线程分别计算对应外积矩阵并将结果累加最终形成完整的_R_i_T_G_i_-相关结果这一操作的关键特点在于每个线程处理得到的外积矩阵尺寸仅为4×M（其中M代表嵌入网络输出的空间维度）因此这些中间结果可以在GPU共享内存中高效地进行求和运算从而避免了完整存储整个_G_i_-相关矩阵所带来的内存消耗问题

3.4.3 冗余0消除

DeePMD-kit 的基准实现在生成每个原子的具体环境矩阵 R_i 时遵循固定的最大邻居表长度 N_m 进行了零填充操作处理。当模型经过压缩优化后, G_i 的值不再由 GEMM 矩阵乘法运算得到,而是采用了并行查找预存表格结合多项式插值运算的方法快速求解出来,从而使得可以根据当前系统中实际存在的邻居数量灵活配置参数设置,避免由于 R_i 矩阵中存在零填充引起的冗余运算带来的额外开销,最终有效地降低了整个推断过程所需的算力资源消耗量

用于训练适用于常压和高压条件的模型时，在执行常压条件下的推断任务时通常能够获得更高的冗余率，在这种情况下采用冗余0优化策略将能够显著减少内存消耗并提升计算效率。

4 典型应用（第一性原理精度分子动力学）测试结果

4.1 RLEKF 优化器测试结果

我们在13种材料体系上进行了系统性比较RLEKF与Adam的性能表现图5通过对比分析可以看出在模型收敛后RLEKF相比Adam显示出更高的加速比结果对比结果显示在其他12种材料上均展现出明显更快的收敛速度

图 5

图****5RLEKF在13种体系下达到收敛时与Adam**的加速比**

The improvement ratios of RLEKF compared to ADAM upon achieving convergence in 13 systems.

我们对RLEKF的预测效果进行了评估。图6详细展示了RLEKF在银元体系、氯化钠体系以及水体系等多种体系下的预测结果与DFT计算值之间的对比关系。其中图中各个样本点与对角线越接近，则表明该模型对该样本点的预测越精确。

图 6

图****6RLEKF在四个典型体系上的预测结果与DFT**计算结果的对比**

Fig.6 Analysis of results from the RLEKF method for four representative systems using DFT calculations

在一定程度上维持了较高的精度水平的RLEKF优化器相比传统Adam方法，在采用大批量版本进行训练时显著提升了训练效率。研究表明，在国产深度计算单元上实施这种大规模并行优化策略后可使DeePMD-kit程序从1天的运行时长缩短至仅需不到4分钟完成整个计算过程

图 7

图****7RLEKF在国产DCU**上的训练加速效果**

Fig.7The training accelerating effect of RLEKF on domestic DCU****

4.2 ALKPU 主动学习平台测试结果

我们在铜体系上展开了系统性评估主动学习平台ALKPU的功能特性。实验结果表明KPU能够准确评估模型对未知构型预测的不确定性而ALKPU采样平台则能精准筛选出具有代表性的构型从而为后续优化工作提供了可靠依据

ALKPU在计算资源消耗方面具有较低的消耗水平。相较于基于多模型委员会的方法而言，在计算资源消耗方面具有明显优势。此外，在采用了该算法所特有的高效优化机制设计之后，并结合大规模并行训练策略的应用下，则能够在几分钟时间内完成整个训练过程

为了验证ALKPU平台的有效性，在600K温度下的铜体系相空间中进行了7个主动学习周期。每个周期结束后，在由 AIMD 轨迹构建形成的验证集中执行测试，并统计整体预测误差。如图8所示，在主动学习周期数量增加的情况下，新生成并被标注的数据点逐步加入训练集合，并导致整体预测误差逐步下降

图 8

图8预测误差随主动学习轮次增加的收敛趋势

Fig.8the pattern of error reduction in prediction errors follows an increasing trend as the number of active learning iterations grows

为了验证ALKPU平台构型采样的有效性，在主动学习循环全部完成后利用收集的数据对模型进行了训练，并将其应用于前述AIMD轨迹验证集进行评估。研究发现，在所有时间步上的预测误差均满足要求（如图9所示）。通过主动学习策略实施采样过程后发现，在目标体系的关键构型分布范围内获得的数据样本能够充分反映其潜在能量场的变化特征。

图 9

图9主动学习完成后模型在全部验证样本上的预测误差

Figure 9 Predictive error rates across all validation sets following completion of active learning

本研究采用上述模型对铜体系进行了分子动力学（MD）模拟，并系统地评估了其基质物理性质。结果与密度泛函理论（DFT）计算结果极为接近。具体数据可见于图1。其中V₀(ų/atom)代表原子平衡体积（Equilibrium volume per atom），EVf代表空位形成能（Vacancy formation energy），而a₀则表示系统处于平衡体积状态下的晶格常数。

表****1****基础物理量计算

Table 1 Calculation of fundamental physical quantities

物理量	V0(Å3/ atom )V0(Å3/ atom )	EVf****(eV)	a0(Å)
DFT	41.6892	0.0002	3.4674
ALKPU	41.2217	0.0016	3.4544

新窗口打开**|下载CSV**

4.3 DeePMD-kit 推理测试结果

在不同配置下（包括V100与A100）GPU环境下对经过优化后的DeePMD-kit单卡性能展开了系统性测试与评估（如图10所示）。所涉及的物理体系包括铜系和水系体系，在V100单卡运行中相比基准版本实现了一定程度的性能优化：对于水系体系而言，在每步原子解算时间上从9.47 μs降至2.58 μs（约3.67倍加速比）；而对于铜系体系则从27.81 μs降至2.87 μs（约9.69倍加速比）。基于更高性能的A100架构，则分别获得了相应的加速比提升：水系体系实现6.91倍加速比、铜系体系则达到显著更高的15.98倍加速比

图 10

图****10DeePMD-kit****单卡推理性能

Fig.10Inference performance of DeePMD-kit on a single GPU****

我们对具有铜体系结构的材料在27,360块V100显卡上进行了大规模推理运算实验研究，并获得了令人满意的实验结果。

图 11

图****11DeePMD-kit****大规模推理性能提升

Fig.11Performance improvements of DeePMD-kit in large-scale inference****

测试结果表明，在采用分段五阶多项式模型压缩基础上进行推理优化的过程中

5 总结与展望

本研究开发了一款融合HPC与AI技术的第一性原理科学智能计算平台。该平台系统性地介绍了第一性原理计算软件与科学智能模型、科学智能计算的关键环节以及科学数据管理流程，并重点阐述了其在以下关键领域的工作：首先是基于科学研究的数据生成机制及数据集构建方法；其次是应用构型空间探索技术并结合主动学习策略；再次是针对高效训练阶段所采用的核心算法创新；最后是对大规模推理过程中的性能优化措施等。目前该平台已完成核心技术系统的全面集成，并已在几个典型物理系统中进行了初步测试验证。下一步的重点将是扩大测试范围并完善相关数据支持；计划于年底前推出内测版供用户评估使用。

科学智能计算从研究范式和工作流层面本质地革新了传统科学计算格局，在突破"维数灾难"的同时为第一性原理计算拓展到更为复杂的现实应用场景提供了新机遇，并预示着技术路径及工作方式将发生根本性变革。
在这一充满活力的发展阶段中，
可以看出大规模、高质量的第一性原理模拟数据对于推动科学智能计算发展的促进作用异常显著，
而在此背景下，
人工智能技术的发展趋势与高效计算方法的应用场景都将面临前所未有的挑战，
成为未来研究的重点方向之一。