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

摘要

【目的】 科学智能（AI for Science）方法正在深刻地改变当前科学计算的格局。其融合了物理模型、人工智能与高性能计算，针对传统科学计算中的高维问题，通过数据拟合的方式实现成量级的增加高精度科学计算问题的时间和空间尺度，正在推动一场科研范式的变革。【方法】 本文针对第一性原理精度的分子动力学，提出一种HPC+AI驱动的科学智能计算平台，针对科学智能在工作流上带来的变化与挑战，从科学数据的生成与数据集制备、构型空间探索与训练样本标注、科学智能模型的高效训练及大规模高效推理等四个方面阐述构建科学智能计算平台的关键技术与流程。【结果】 本文所提出的计算平台在整合科学智能计算工作流的基础上，针对HPC+AI驱动的第一性原理精度分子动力学这一典型应用，提出了基于卡尔曼滤波的主动学习策略；改进了拟二阶AI模型训练方法，实现训练时间从天到分钟级的加速；利用五阶多项式AI模型压缩技术实现在同等硬件条件下模型推理的体系规模提高1个数量级，到解时间提高3-9倍。【结论】 通过上述工作的整合，形成一套可用于第一性原理精度分子动力学计算的科学智能计算平台。【局限与展望】 科学智能计算方法与工作流仍处于蓬勃发展阶段，在高精度数据、更通用AI模型和高效的计算方法等方面仍面临巨大的挑战，也将成为本文工作在未来的重要探索方向。

关键词： 科学智能; 第一性原理计算; 分子动力学; 主动学习; 卡尔曼滤波; 模型压缩

引言

第一性原理计算依据量子力学的基本理论，可以实现从底层物理规律出发定量地研究复杂问题，材料、化学、生物等领域的众多科学问题均可通过第一性原理计算进行求解。然而，第一性原理计算的复杂度高，其计算量随着体系规模呈指数增长，当体系规模较大时便会遇到“维数灾难”，极大地限制了其所能实际模拟的时间和空间尺度：即使利用目前最先进的超级计算机，也只能针对数千原子体系进行皮秒量级的第一性原理模拟。而领域科学家所关心的大量复杂现实科学问题，如薄膜生长、表面催化、材料设计以及蛋白质折叠等都需要对大体系进行长时间的模拟。如何提高第一性原理计算所能模拟的时间和空间尺度，使其能够应用于解决众多物理、化学、生物领域的复杂现实科学问题，已经成为第一性原理计算领域发展的一个关键问题。

以机器学习为代表的人工智能方法为解决传统科学计算所面临的维数灾难提供了新的契机。相关研究表明，神经网络能够以接近蒙特卡洛方法[1]的复杂度来有效地表示或者逼近高维空间的函数[2-3]，这为应用神经网络来拟合科学计算中的高维问题提供了理论基础，而神经网络在一系列AI for Science应用中的成功实践[4⇓⇓⇓⇓-9]也证实了其在解决科学问题中的潜力。在第一性原理计算领域，神经网络模型自2007年起就开始得到重视和逐步应用，截至2021年，以DeePMD-kit为典型代表的科学智能计算模型通过结合物理模型、人工智能方法和高性能计算，利用超级计算机实现了上亿原子体系的纳秒级模拟[4,10]，使第一性原理方法朝着解决复杂现实科学问题的方向又迈进了一步。

尽管如此，第一性原理计算在可模拟的时间和空间尺度上，距离直接解决材料、化学、生物等领域的复杂现实科学问题仍有若干数量级的差距。要填补这个差距，一方面有赖于领域专用超级计算机等超算体系结构的进一步发展，另一方面仍需要在科学应用以及科学智能计算方法上继续突破。科学智能计算方法从研究范式和工作流层面深刻地改变了传统科学计算的格局，也带来了新的问题与挑战：其首要体现在于科学数据的使用方式发生了根本性的转变，如何高效地生成以及有效地利用科学数据成为倍受关注的重要问题；其次，新的工作流也对训练样本的使用效率、AI模型的泛化与外推能力、模型训练时间以及响应快速迭代需求的能力等方面提出了新的挑战。

本文面向第一性原理计算领域科学智能方法带来的变化与挑战，提出一种基于HPC+AI的科学智能计算平台，分别从第一性原理模拟数据的高效生成与数据集制备、面向应用的构型空间探索与主动学习、科学智能模型的高效训练及大规模高效推理等四个方面提出优化策略。在整合科学智能计算工作流的基础上，针对HPC+AI驱动的第一性原理精度分子动力学这一典型应用，提出基于卡尔曼滤波的主动学习策略，改进科学智能模型的拟二阶训练方法，对科学智能模型进行五阶多项式模型压缩与大规模推理优化。通过上述工作流整合、关键技术优化与平台化集成，进一步提高科学智能方法的应用范围、技术性能并形成能够满足现实应用需求的科研生产力平台。

1 相关工作

Behler[11]于2007年首次提出通过神经网络来拟合第一性原理精度的高维原子势能面函数，通过将体系的总势能分解为体系中每个原子的势能之和，同时利用“近视原理”来对原子势能的计算进行近似处理，Behler提出的模型成功实现了对目标体系的任意规模高维原子势能面的灵活拟合，为AI驱动的第一性原理精度分子动力学走向实用奠定了基础。之后，伴随着人工智能和深度学习相关技术的爆发，这一领域得到了快速的发展。

2017年，Schütt等人提出基于连续滤波卷积的SchNet[12]模型，该模型通过引入光滑的滤波函数来使卷积层能够以非网格化的方式处理连续的原子位置变化，在实现平移不变性和旋转不变性的同时生成光滑的拟合势能面，保证了根据原子势能面计算出的力场符合能量守恒定律，实现了化学精度的原子势能面和原子力场拟合。2018年，张林峰等人提出端到端且保证物理对称性的原子间势能面模型[13]，该模型通过输入的原子位置自动生成保证平移不变性、旋转不变性以及交换不变性的特征来实现高效的端到端拟合。2020年DeePMD-kit[6]软件在超级计算机上实现上亿原子体系1纳秒/天的模拟速度[4]，推动了第一性原理精度分子动力学的跨越式发展。

在分子建模领域，图神经网络（Graph Neural Networks, GNN）以其网络结构与分子结构的天然相似性获得广泛青睐，也涌现出一批支持第一性原理精度的应用。DimeNet++[14]利用GNN的消息传递机制，通过原子对的距离信息和原子三联体的角度信息来预测分子的量子力学性质并取得了先进的结果；GraphNVP[15]结合变分自编码器（Variational Autoencoder, VAE）和GNN来学习化学分子的分布；CGCNN[16]基于GNN专门用于对晶体材料进行建模和预测。

近年来，以Transformer[17]为基础的大模型技术在自然语言处理及图像处理领域取得了瞩目的突破性进展。大模型技术在第一性原理精度分子动力学模拟中也得到了初步应用并取得一定进展。2021年，微软亚洲研究院团队提出Graphormer[18]模型。传统GNN分子建模方法当网络层数较多时会出现过度平滑问题，限制了网络规模以及模型的表达能力，Graphormer模型通过将图的结构信息编码到Transformer的自注意力层中成功突破了这一限制。2022年，基于Transformer架构的原子间势能面大模型DPA-1[7]，通过在深度势能模型的基础上引入元素类型编码，同时基于Transformer实现类比多体相互作用的原子间信息交互，最终实现了涵盖多种元素的预训练大模型并完成在多种下游任务上的迁移学习，将“大数据预训练+少量具体任务数据微调”的模式引入到了第一性原理精度分子动力学领域。

在科学数据集方面，依托国家“材料基因工程”发展战略和大数据驱动的科研创新趋势，国内科研机构和大学也建立了相关的领域数据库，为第一性原理领域的科学智能计算提供了有力支撑。例如由中国科学院物理研究所刘淼博士等人建立的Atomly[19]材料科学数据库提供了30余万种无机化合物的数据；北京大学潘峰教授团队也创建了高精度材料科学数据库。由北京科学智能研究院(AISI)等单位建立的科学智能广场 (AIS-Square)[20]是科学智能数据、模型、工作流的共享开源协作平台，提供涵盖了金属、半导体、氧化物和过渡金属化合物的多种材料计算数据，同时还提供了50多个特定场景的专用科学智能模型及配套的工作流。中国科学院计算机网络信息中心自主研发的科学数据银行（Science Data Bank）[21]是一个通用型论文关联数据存储平台，汇集了各领域的科学数据，数据体量达到342,819+ GB，不仅提升了科研数据成果的价值，也促进了科学研究和教育的发展。

科学数据存在获取成本高、覆盖范围不够广泛的特点，因此，能够对应用的构型空间进行有效探索、对训练样本进行精准标注的主动学习方法就成为了高效使用科学数据的关键技术之一。Frederiksen等人[22]最先提出将主动学习方法应用于经验PES（Potential Energy Surface）模型，他们提出了贝叶斯委员会查询策略，该方法可以评估不同类型的原子间势，也可用于优化拟合数据库。Botu等人[23]提出一种根据给定原子构型与训练集构型之间的距离来评估预测误差的方法。还有一系列工作基于高斯过程回归力场[24⇓⇓-27]，通过贝叶斯推断的方差来评估不确定性。Podryabinkin等人[28]提出了一种基于D-最优标准的主动学习方法，该方法能够减少确定参数的不确定性，适用于所有线性势，即能量线性依赖于参数的势。当前学术界和工业界广泛使用的主动学习采样平台DP-GEN[8,29⇓⇓-32]通过多个结构相同但初始化训练参数不同的模型进行集合预测的方法来度量不确定性，该平台已经支持了一系列应用研究工作[30⇓⇓⇓-34]