基于神經網絡的變分蒙特卡羅（NN-VMC）已成為一種有前途、有潛力的從頭計算量子化學的尖端技術。然而，現(xiàn)有方法的高計算成本，阻礙了它們在現(xiàn)實化學問題中的應用。

字節(jié)跳動與北京大學團隊介紹了 NN-VMC 的最新研究進展，它實現(xiàn)了顯著的加速率，從而極大地將 NN-VMC 的適用性擴展到更大的系統(tǒng)。

該團隊的關鍵設計是一個名為 Forward Laplacian 的計算框架，它通過有效的前向傳播（forward）過程來計算與神經網絡相關的 Laplacian（這是 NN-VMC 的瓶頸）。

然后，研究人員證明 Forward Laplacian 可以進一步促進各個方面的加速方法的更多發(fā)展，包括稀疏導數(shù)矩陣的優(yōu)化和高效的網絡設計。

該團隊的方法使 NN-VMC 能夠研究更廣泛的系統(tǒng)，為其他從頭算方法提供有價值的參考。

該研究以「A computational framework for neural network-based variational Monte Carlo with Forward Laplacian」為題，于 2024 年 2 月 13 日發(fā)布在《Nature Machine Intelligence》。

準確求解與時間無關的電子薛定諤方程，可以得出給定量子力學系統(tǒng)的基本屬性。量子蒙特卡羅（QMC）是求解薛定諤方程最重要的從頭算方法之一，廣泛應用于量子化學的各種場景。

然而，在 QMC 中，解決方案的準確性在很大程度上取決于 ansatz 的選擇，這需要大量的專業(yè)知識。與其他確定性方法相比，例如耦合簇單簇和擾動三元組的雙簇 [CCSD(T)] ，這種限制阻礙了 QMC 的適應性，后者被廣泛認為是「黃金標準」。

近期，深度學習徹底改變了量子化學領域，使科學家可以獲得薛定諤方程更準確的解。開創(chuàng)性的方法之一是基于神經網絡的實空間變分蒙特卡羅（NN-VMC），例如 FermiNet 和 PauliNet。通過變分原理可以得到基態(tài)波函數(shù)，從而使能量的期望值最小化。基于這一事實，NN-VMC 方法利用深度神經網絡對波函數(shù)進行參數(shù)化，并以能量作為損失函數(shù)來優(yōu)化網絡參數(shù)。

受益于神經網絡的卓越能力，NN-VMC 方法在實現(xiàn)不同系統(tǒng)的化學準確性方面表現(xiàn)出了穩(wěn)健的結果。然而，需要注意的是，這些方法在模型訓練過程中通常會產生大量的計算成本。例如，在現(xiàn)代硬件上學習苯二聚體系統(tǒng)的波函數(shù)需要大約 10,000 個圖形處理單元 (GPU) 小時，這使得擴展到更大的系統(tǒng)具有挑戰(zhàn)性。

在最新的工作中，字節(jié)跳動和北京大學的研究人團隊解決了 NN-VMC 方法中計算效率的挑戰(zhàn)，特別是對于大規(guī)模系統(tǒng)。

NN-VMC 需要計算神經網絡輸入的 Laplacian 以獲得損失；該項的計算消耗了整個訓練時間的很大一部分，并成為學習過程中的主要瓶頸。

具體來說，之前的工作首先使用深度學習工具包中的自微分（AutoDiff）方法計算 Hessian 矩陣，然后通過跡推導 Laplacian 矩陣。這樣的過程需要執(zhí)行昂貴的導數(shù)傳播，顯著降低了 NN-VMC 的整體訓練速度。

為了觀察這個問題，研究人員首先開發(fā)了一個新的計算框架，名為 Forward Laplacian。?

與從 Hessian 間接導出 Laplacian 的常用方法相比，F(xiàn)orward Laplacian 通過精心設計的前向傳播過程直接計算值；研究人員在數(shù)學上證明這種方法效率更高，因為它消除了不必要的計算和傳播。

其次，研究人員證明這種計算方式不僅加速了過程，而且為開發(fā) NN-VMC 先進技術鋪平了道路。在與 NN-VMC ansatz 相關的 Forward Laplacian 中，許多中間導數(shù)表現(xiàn)出稀疏性并且可以在相當大的程度上進行優(yōu)化。

該團隊還設計了一種名為 LapNet 的高效神經網絡架構，它可以使用精心設計的具有稀疏導數(shù)的注意力塊來更好地利用 Forward Laplacian方法的優(yōu)勢。總之，這些進展使研究人員能夠研究超出現(xiàn)有 NN-VMC 軟件包能力的原子、分子和化學反應。

該團隊在計算絕對能量跟隨方面在各種系統(tǒng)上評估他們的方法。所有結果一致表明，LapNet 與 Forward Laplacian 方法相結合，可以獲得準確的能量估計，同時顯著降低模型訓練的計算成本。

鑒于這些穩(wěn)健的結果，研究人員進一步探索該方法是否可以在不同的實際場景中學習更有用的量（相對能量），包括化學反應的勢壘、過渡金屬的電離能和分子之間的非共價相互作用。

結果表明，使用 NN-VMC 方法獲得的相對能量與使用黃金標準計算方法和實驗結果獲得的相對能量一致，這表明使用深度學習來解決量子力學系統(tǒng)的巨大潛力。

雖然該團隊提出的方法在大多數(shù)情況下，可以產生與金標準或實驗結果密切相關的相對能量；但在少數(shù)情況下，該方法輸出的結果與現(xiàn)實情況之間仍然存在差異。

該團隊認為這種不一致是因為現(xiàn)有的 NN-VMC 方法沒有充分整合所有關鍵的化學和物理知識。研究人員相信，將適當?shù)幕瘜W和物理知識（例如尺寸擴展性和尺寸一致性）編碼到神經網絡中對于未來推進 NN-VMC 技術至關重要。

雖然，該研究的主要目標是解決與 NN-VMC 中 Laplacian 計算相關的計算瓶頸，但是 Forward Laplacian 方法的預期用途可以擴展到量子力學領域之外的各種場景。

例如，F(xiàn)orward Laplacian 方法在計算 Laplacian 時可以代替以前的方法，并且提高了效率。因此，它可以加速其他基于神經網絡的偏微分方程求解器。

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