DeePKS-kit:天梯亦可上,只要肯登攀

化学

机器学习帮助我们攀登密度泛函中的 Jacob 天梯

在电子结构性质的计算中,我们通常需要根据精度与效率的权衡来选择合适的方法。耦合簇方法(如 CCSD、CCSD(T))被认为具有化学精度,但它的计算复杂度在 O(N7)O(N^7) 量级;另一方面,基于 Kohn-Sham 理论的各种密度泛函方法的计算复杂度可以控制在 O(N3)O(N4)O(N^3)\sim O(N^4),但它们的精度常常不能令人满意。

我们注意到,密度泛函方法虽然由于平均场近似而精度受限,但仍然在量子化学与凝聚态物理研究中占据了主流地位。这是由于,密度泛函模型的核心——将体系能量表示为电子密度的泛函——具有很高的灵活性,人们在实践中往往能够根据物理直觉和经验数据的结合来调整密度泛函模型,从而达到所需的精度。也正因为如此,构建既准确又高效的密度泛函模型事实上成为了一个参数拟合问题——全宇宙最有挑战性也最激动人心的参数拟合问题之一。

从一个基础的密度泛函模型出发逐步改进、并最终到达化学精度的过程通常被形象地比喻为攀登「Jacob 天梯」的过程。在《创世纪》中,犹太人的祖先 Jacob 梦见一架通天的梯子,他登上梯子到达了天堂,并受到了神的祝福。而在科学家的手中,通过向泛函中不断加入更多的考虑因素,就能从 Hartree-Fock 精度的人间到达化学精度的天堂。然而长期以来,人们缺乏适合的工具来系统地改善密度泛函模型,依赖少量数据的手工参数调整不仅费时费力,还常常只能应用于特定的体系,不具有很好的可迁移性。总而言之,攀登「Jacob 天梯」的每一步都困难重重。

密度泛函模型的「Jacob 天梯」

现在,机器学习与物理建模的结合正在改变着这一切:前者是处理复杂信息与数据的有效工具;后者是人们认知物理世界的科学方式。机器学习帮助我们以新的视角来重新审视数据驱动的密度泛函模型构建,将密度泛函模型的通用性问题归结为数据集的表示能力和模型的泛化能力;机器学习还提供了一系列神经网络架构和优化算法来参数化和改进密度泛函模型。因此,我们有望借助机器学习的力量,改变以往构建密度泛函模型的范式,让「Jacob 天梯」的攀登变得前所未有地简单。

方法回顾

在这一理念的引导下,2020 年 8 月,我们发表了论文 DeePKS: a comprehensive data-driven approach towards accurate density functional theory,提出了构建准确且高效的密度泛函模型的通用机器学习框架,并且利用这一框架训练了具有化学精度的密度泛函模型。DeePKS 的设计依据以下几个原则:必须能够适用于各种类型的体系;必须能用小体系上训练的模型计算大体系;必须能充分利用现有的多种电子结构数据,包括在给定构型下的能量、力和电子密度等。为此,我们将体系的总能量表示为一种基准泛函模型给出的能量加上一个由机器学习给出的修正项:

E=min{φi},φiφj=δij{Ebase[{φi}]+Eδ[{φiω}]}E=\min_{\left\{\varphi_{i}\right\},\left\langle\varphi_{i} \mid {\varphi}_{j}\right\rangle=\delta_{i j}}\left\{E_{\rm base}\left[\left\{\varphi_{i}\right\}\right]+E_{\delta}\left[\left\{\varphi_{i} \mid \omega\right\}\right]\right\}

其中 {φi}\{\varphi_i\} 是一组 Kohn-Sham 轨道,EbaseE_{\rm base} 是基准模型(例如 Hartree-Fock 或 PBE 泛函),ω\omega 是修正项 EδE_\delta 中的参数。而在修正项的计算过程中,我们首先将密度矩阵预处理为每个原子 II 上的描述符 dI\mathbf d^I,然后通过神经网络 FNN\mathcal F^{\rm NN} 计算出

Eδ=IFNN(dIω)E_\delta=\sum_I\mathcal F^{\rm NN}(\mathbf d^I\mid\omega)

在构建好模型之后,我们将模型的优化问题表述为在自洽场方程约束下最小化各种电子结构性质 Y[{φi}ω]Y[\{\varphi_i\}\mid\omega] 和高精度标签 YlabelY_{\rm label} 的均方误差的过程。为了展示这样的模型的泛化能力,我们在 CCSD(T) 精度的、包含 7212 个不同分子的 QM7b-T 数据库上训练,测试误差仅为 0.2 kcal/mol,在化学精度所要求的 1 kcal/mol 以内。

DeePKS 方法在 QM7b-T 数据库上的表现

软件介绍

不仅如此,我们还意识到,目前电子结构计算社区急需能够将机器学习平台和电子结构计算软件桥接的开源工具,帮助更多的研究者提升开发效率、改善算法的可重现性。这样的工具不仅应当为用户提供易于使用的界面,还应当为研究者提供开发新算法的灵活的接口。因此,我们在 GitHub 上以 LGPL-3.0 协议开源了 DeePKS-kit,它以机器学习框架 PyTorch 为前端、量子化学软件 PySCF 为后端打造了直观的机器学习密度泛函模型的构建工作流。

DeePKS-kit 工作流程

上图展示了 DeePKS-kit 的工作流程。我们首先用基准泛函模型在 PySCF 中求解一个给定的体系,生成相应的 Kohn-Sham 轨道和描述符;然后,我们在 PyTorch 中初始化一个模型,读入描述符和标签数据,优化模型中的参数;优化若干步之后,再将模型传递给 PySCF,用新的泛函模型再次求解体系得到更新的描述符;如此循环往复直到泛函模型达到所需精度。

DeePKS-kit 通过命令行接口 deepks train deepks scf 提供了上述每一步的功能,还提供了 deepks iterate 来将上述过程自动化,包括管理计算过程中生成的各种文件、调度计算任务、测试训练过程中初步得到的模型、以及向高性能计算集群申请计算资源等。工作流程的标准化不仅方便了新算法的快速验证,更使得后续大规模的模型构建以及投入应用成为了可能。

未来计划

当然,DeePKS-kit 还远不止于此——作为 Deep Modeling 社区的重要组成部分之一,DeePKS-kit 正与很多其他项目产生着丰富的联系:

  • 我们正在积极推动 DeePKS-kit 与开源电子结构计算软件的结合,为有限体系和凝聚态体系准确求解提供更多可能性;
  • DP-Dispatcher 正在积极开发,它是 Deep Modeling 社区的下一代调度工具,将为 DeePKS-kit 提供更完善的工作流程定义;
  • DeePKS-kit 本身也将不断完善,支持对更多的(不限于能量、力和电子密度)标签的学习,进一步提高模型的表现。

总而言之,DeePKS-kit、ABACUS 与其他上下游配套工具链将是 Deep Modeling 社区将建模尺度延伸至电子结构的重要里程碑。

如何加入?

欢迎你在 DeePKS-kit 仓库中撰写 Issues,发起 Discussions,甚至提交 Pull Requests。具体来说:

  • 如果你正好也在探索数据驱动的密度泛函模型的构建方式,你可以与我们进行深入的学术讨论,丰富完善 DeePKS-kit 的训练模块;
  • 如果你是熟悉电子结构计算软件的硬核开发者,希望为你的软件加入更加酷炫的功能,你可以了解我们的前后端分离的设计,尝试把 DeePKS-kit 训练出的密度泛函模型集成到你的软件之中;
  • 如果你正致力于量子化学和凝聚态物理的研究,希望体验更强大的计算工具,你可以成为我们的第一批用户,对我们上线的功能进行尝鲜,并从实际的需求出发给我们提出宝贵建议,推动 DeePKS-kit 变得更好。

结语

在高性能计算技术与机器学习方法飞速发展的今天,我们看到了电子结构计算的新的可能性。以往,很多商用的电子结构计算软件不易扩展,使用的方法落后于前沿的科学进展;新的方法的影响力往往受限于特定的几个学术课题组,难以进一步扩大用户群体;密度泛函方法经过 60 年的发展积累了大量可用的模型,但在具体问题的建模中,选用什么模型有时还需要大量经验的积累。DeePKS-kit 带领 Deep Modeling 社区中的周边设施将软件、模型和应用整合到一个框架下,对密度泛函方法发展的新范式做了有益的探索。

天梯亦可上,只要肯登攀!


站内评论

邮箱仅用于回复提醒,不会公开展示


站内评论加载中……


站外评论

站外评论加载中……