加速量子化学计算，字节&北大量子蒙特卡罗研究进展登Nature子刊

编辑 | 萝卜皮

基于神经网络的变分蒙特卡罗（NN-VMC）已成为一种有前途、有潜力的从头计算量子化学的尖端技术。然而，现有方法的高计算成本，阻碍了它们在现实化学问题中的应用。

字节跳动与北京大学团队介绍了 NN-VMC 的最新研究进展，它实现了显著的加速率，从而极大地将 NN-VMC 的适用性扩展到更大的系统。

该团队的关键设计是一个名为 Forward Laplacian 的计算框架，它通过有效的前向传播（forward）过程来计算与神经网络相关的 Laplacian（这是 NN-VMC 的瓶颈）。

然后，研究人员证明Forward Laplacian 可以进一步促进各个方面的加速方法的更多发展，包括稀疏导数矩阵的优化和高效的网络设计。

该团队的方法使 NN-VMC 能够研究更广泛的系统，为其他从头算方法提供有价值的参考。

该研究以「A computational framework for neural network-based variational Monte Carlo with Forward Laplacian」为题，于 2024 年 2 月 13 日发布在《Nature Machine Intelligence》。

准确求解与时间无关的电子薛定谔方程，可以得出给定量子力学系统的基本属性。量子蒙特卡罗（QMC）是求解薛定谔方程最重要的从头算方法之一，广泛应用于量子化学的各种场景。

然而，在 QMC 中，解决方案的准确性在很大程度上取决于 ansatz 的选择，这需要大量的专业知识。与其他确定性方法相比，例如耦合簇单簇和扰动三元组的双簇 [CCSD(T)] ，这种限制阻碍了 QMC 的适应性，后者被广泛认为是「黄金标准」。

近期，深度学习彻底改变了量子化学领域，使科学家可以获得薛定谔方程更准确的解。开创性的方法之一是基于神经网络的实空间变分蒙特卡罗（NN-VMC），例如 FermiNet 和 PauliNet。通过变分原理可以得到基态波函数，从而使能量的期望值最小化。基于这一事实，NN-VMC 方法利用深度神经网络对波函数进行参数化，并以能量作为损失函数来优化网络参数。

受益于神经网络的卓越能力，NN-VMC 方法在实现不同系统的化学准确性方面表现出了稳健的结果。然而，需要注意的是，这些方法在模型训练过程中通常会产生大量的计算成本。例如，在现代硬件上学习苯二聚体系统的波函数需要大约 10,000 个图形处理单元 (GPU) 小时，这使得扩展到更大的系统具有挑战性。

在最新的工作中，字节跳动和北京大学的研究人团队解决了 NN-VMC 方法中计算效率的挑战，特别是对于大规模系统。

图示：计算过程示意图。（来源：论文）

NN-VMC 需要计算神经网络输入的 Laplacian 以获得损失；该项的计算消耗了整个训练时间的很大一部分，并成为学习过程中的主要瓶颈。

具体来说，之前的工作首先使用深度学习工具包中的自微分（AutoDiff）方法计算 Hessian 矩阵，然后通过迹推导 Laplacian 矩阵。这样的过程需要执行昂贵的导数传播，显著降低了 NN-VMC 的整体训练速度。

为了观察这个问题，研究人员首先开发了一个新的计算框架，名为 Forward Laplacian。

与从 Hessian 间接导出 Laplacian 的常用方法相比，Forward Laplacian 通过精心设计的前向传播过程直接计算值；研究人员在数学上证明这种方法效率更高，因为它消除了不必要的计算和传播。

其次，研究人员证明这种计算方式不仅加速了过程，而且为开发 NN-VMC 先进技术铺平了道路。在与 NN-VMC ansatz 相关的 Forward Laplacian 中，许多中间导数表现出稀疏性并且可以在相当大的程度上进行优化。

图示：LapNet 架构。（来源：论文）

该团队还设计了一种名为 LapNet 的高效神经网络架构，它可以使用精心设计的具有稀疏导数的注意力块来更好地利用 Forward Laplacian方法的优势。总之，这些进展使研究人员能够研究超出现有 NN-VMC 软件包能力的原子、分子和化学反应。

该团队在计算绝对能量跟随方面在各种系统上评估他们的方法。所有结果一致表明，LapNet 与 Forward Laplacian 方法相结合，可以获得准确的能量估计，同时显著降低模型训练的计算成本。

图示：不同 NN-VMC 方法的效率和性能比较。（来源：论文）

鉴于这些稳健的结果，研究人员进一步探索该方法是否可以在不同的实际场景中学习更有用的量（相对能量），包括化学反应的势垒、过渡金属的电离能和分子之间的非共价相互作用。

图示：系统规模的缩放效应。（来源：论文）

图示：使用 LapNet 进行相对能量估计。（来源：论文）

结果表明，使用 NN-VMC 方法获得的相对能量与使用黄金标准计算方法和实验结果获得的相对能量一致，这表明使用深度学习来解决量子力学系统的巨大潜力。

虽然该团队提出的方法在大多数情况下，可以产生与金标准或实验结果密切相关的相对能量；但在少数情况下，该方法输出的结果与现实情况之间仍然存在差异。

该团队认为这种不一致是因为现有的 NN-VMC 方法没有充分整合所有关键的化学和物理知识。研究人员相信，将适当的化学和物理知识（例如尺寸扩展性和尺寸一致性）编码到神经网络中对于未来推进 NN-VMC 技术至关重要。

虽然，该研究的主要目标是解决与 NN-VMC 中 Laplacian 计算相关的计算瓶颈，但是 Forward Laplacian 方法的预期用途可以扩展到量子力学领域之外的各种场景。

例如，Forward Laplacian 方法在计算 Laplacian 时可以代替以前的方法，并且提高了效率。因此，它可以加速其他基于神经网络的偏微分方程求解器。

论文链接：https://www.nature.com/articles/s42256-024-00794-x