编辑 | 萝卜皮
鉴于神经网络的变分蒙特卡罗（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