从 DFT 先去 DeePKS 再到 DeePMD | DeePKS基础篇

💭目录

• DP模型的开发流程
• Kohn−Sham简介
• DeePKS模型的开发流程
  • 第一手👋：描述符
  • 第二手👋：损失函数
  • 第三手👋：在训练集中的DeePKS模型迭代
  • 第四手👋：在测试集中的DeePKS模型计算
• 总结

DP模型的开发流程

近年来，机器学习势函数 (下面都以DP为例)的发展推动了分子模拟领域的研究，使得针对大体系的高精度分子模拟成为可能。然而，机器学习势函数的训练首先需要生成大量的DFT数据，对于精度要求高的场景（如杂化泛函），生成这些数据会耗费大量计算资源。DeePKS 模型的研发旨在能够有效缓解这一问题，借助AI构建起DFT与DPMD之间的桥梁，即[从 DFT 先去 DeePKS 再到 DeePMD]。

为了构建高质量的DP模型的训练数据集，我们可以使用DP-GEN(Deep Potential GENerator) 软件。DP-GEN是一个自动化的训练数据集生成工具，分为MD构象采样、DFT标定数据、训练DP模型三个步骤。这三个步骤可以通俗的理解为：在中学备考时代，我们会收集有价值的错题、然后标注参考答案、不断巩固训练大脑完成从题目到答案的映射。

具体而言：

我们的关注点在于DFT标定数据这一步骤。在量子化学程序中，比如我们想研究有机体系，气相反应计算往往以杂化泛函起步 (下面以杂化泛函 meta-GGA M06-2X 为例)。但这对于周期性体系液相反应计算来讲，杂化泛函过于昂贵，大部分文章会选择用普通泛函 GGA PBE (目前引用次数超17万次，似乎也说明天下苦只能选PBE久矣)。

沿着Jacob天梯向上，如果我们想用双杂化泛函，甚至是目前的黄金标准CCSD(T)，这几乎是不可能完成的挑战！

但DeePKS有望提供一些转机。带入到DP模型的开发流程场景中，原本需要生产上万个M06-2X精度数据作为DP训练集，将由上百个的M06-2X精度计算作为DeePKS训练集。训练后的DeePKS模型将代替M06-2X，再去计算上万个单点来组成DP训练集，最后训练DP模型(精度接近M06-2X)。该过程预计能够节省一个数量级的时间成本。

Kohn−Sham简介

别着急，在介绍DeePKS原理之前，也许您好奇为什么叫DeePKS，DeeP想必不言而喻，这里我们就以Kohn−Sham方程简单过渡一下。故事涉及到两次诺奖，我们只好长话短说。

即使有了薛定谔方程、狄拉克方程，解析求解三体系统问题仍过于复杂，再加激发态则是难上加难。无奈寻求近似方法，目前较常用是Kohn−Sham方程，其借助密度泛函，可通过SCF(self-consistent field)迭代得到近似解，其中的能量和受力主要构成了DP的训练集，DeePKS的训练过程中则额外会利用SCF收敛后的密度矩阵和原子轨道信息。

即便有了Kohn−Sham方程，事情似乎也不简单，如何构造精确的exchange-correlation项仍然是个开放性问题 (这里所列的M06-2X, SCAN, PBE泛函的不同，主要源自exchange-correlation，其他项比较容易确定)。

不过，DeePKS不是去构造exchange-correlation泛函，而是通过机器学习对低精度的泛函进行优化，利用神经网络修正项去学习baseline泛函（低精度、低成本，以PBE为例）与target泛函（高精度、高成本，以M06-2X为例）之间的能量与力的差值。其中，修正项带来的额外计算成本与低精度泛函相当，并远低于高精度泛函。

DeePKS模型的开发流程

第一手👋：描述符

为满足不同大小体系的通用性，类似于BPNN最初将总能量化为原子能量的求和，DeePKS中能量的修正项也具有局域性。修正项作为神经网络的输出，神经网络的输入是描述符$\mathbf{d}$,而DeePKS描述符有以下要求：

如果您对投影密度矩阵$D$的构造有困惑，下图提供了简单但并不十分严谨的例子供您理解。

如果您还是觉得太抽象，并不想了解太多细节，也可以先跳过这一步骤，那就请记住DeePKS描述符包含了坐标、轨道、密度矩阵就够啦。

密度矩阵信息极为丰富，所以DeePKS训练集中的构象数量可以比DP少两个数量级，但仍然包含了大量特征。

第二手👋：损失函数

得到能量修正项$E_{\delta }$后，我们可以获得总能量$E_{\mathrm{DeePKS}}=E_{\mathrm{baseline}}+E_{\delta }$，进而得到总受力${\mathbf{F}}_{\mathrm{DeePKS}}={\mathbf{F}}_{\mathrm{baseline}}+{\mathbf{F}}_{\delta }$。值得一提的是$E_{\delta }$和${\mathbf{F}}_{\delta }$都是由神经网络参数和DFT计算输出决定的。

训练神经网络的过程即最小化DeePKS和target之间的能量、受力差异，最后的拟合目标是$E_{\mathrm{DeePKS}}\approx E_{\mathrm{target}}$和${\mathbf{F}}_{\mathrm{DeePKS}}\approx {\mathbf{F}}_{\mathrm{target}}$。

第三手👋：在训练集中的DeePKS模型迭代

故事本该到这里就结束了，但事情并非如此简单。我们还忽略了一个因素 (想想我们用第一性原理程序分别计算获得$E_{\mathrm{PBE}}$和$E_{M06-2X}$时，SCF迭代过程中有什么不同？据此，读者朋友或许可以先推测一下是少了哪一步呢？)，以至于上述的训练过程并不十分顺利。

我们来揭晓答案。上图根据笔者的理解改编自对应参考文献，发现是需要通过较复杂的DeePKS iteration来确定最终的DeePKS模型,并且用到了哈密顿量修正项${\hat{V}}_{\delta }$，即需要对总哈密顿量进行校正(${\hat{H}}_{\mathrm{DeePKS}}={\hat{H}}_{\mathrm{PBE}}+{\hat{V}}_{\delta }$)。

我们已经有了$E_{\delta }$,为什么还需要${\hat{V}}_{\delta }$呢？原因是如果我们只用基于$E_{\mathrm{PBE}}$计算获得的密度矩阵信息(SCF迭代完后获得)，去预测能量差$E_{M06-2X}-E_{\mathrm{PBE}}$，其实会存在基态信息不匹配的问题，即$initial{\hat{H}}_{\mathrm{DeePKS}}\approx baseline{\hat{H}}_{\mathrm{PBE}}\ne target{\hat{H}}_{M06-2X}$。

这意味着只用基于$E_{\mathrm{PBE}}$计算获得的密度矩阵信息作为主要输入，本身存在一些缺陷，自然加大了拟合的难度。怎么办呢？我们想获取$E_{M06-2X}$的基态信息作为输入，但又不能直接去算... (因为我们的出发点是用DeePKS代替M06-2X)

解决方案是通过DeePKS iteration，即在SCF小循环外再加大循环。小循环的输出作为神经网络校正项的输入，神经网络校正项归属于总哈密顿量，又嵌套在SCF小循环中，基态信息在大循环中不断被修正，最终当loss足够小时，便获得了最后的DeePKS模型，即此时$final{\hat{H}}_{\mathrm{DeePKS}}\approx target{\hat{H}}_{M06-2X}$。对比没有大循环的训练过程，DeePKS iteration能够获得更小的训练误差。

第四手👋：在测试集中的DeePKS模型计算

在训练集中需要大小双循环来训练DeePKS模型，在测试集上推理或者进行上万单点计算时，我们可不想再用耗时的大循环了，这时仅需SCF小循环。

通过上图的${\hat{V}}_{\delta }$公式拆解，我们不难发现其来自于训练好的神经网络参数项，以及PBE DFT小循环计算中的参数。也就意味着我们不是分别算$E_{PBE}$和$E_{\delta }$再求和，只是总能量近似可以拆成两项，方便读者理解才这样书写。

实际上，我们最终获得的是校正后PBE泛函，也可以理解为PBE based DeePKS模型，最终是用其进行SCF从初猜波函到迭代收敛，获得能量和受力用于构造DP训练集，此外一般单点计算能获取的信息如密度矩阵也便能得到。相应的计算时间PBE $\approx$ DeePKS < M06-2X，而精度则是PBE < DeePKS $\approx$ M06-2X。

至此DFT和DPMD之间的桥梁DeePKS在理论上就搭建完毕了。

总结

📖 本Notebook根据笔者的阅读原文献和实践经验，总结了DeePKS基础篇，其借助AI构建起DFT与DPMD之间的桥梁，即[从 DFT 先去 DeePKS 再到 DeePMD]。

🎩 不过我们暂时还不能说DeePKS已经十分完美了，笔者在使用工作中也会遇到大循环loss下降困难、不同体系迁移性还不够好、加入${\hat{V}}_{\delta }$使得SCF迭代难收敛的问题。这些问题可能除了方法本身，当然也涉及到DFT参数设置、数据集构建、超参数选择等诸多因素的影响，期待大家来一起完善！