近年來，基于神經(jīng)網(wǎng)絡(luò)的偏微分方程求解器在各領(lǐng)域均得到了廣泛關(guān)注。其中，量子變分蒙特卡洛方法（NNVMC）在量子化學(xué)領(lǐng)域異軍突起，對于一系列問題的解決展現(xiàn)出超越傳統(tǒng)方法的精確度 [1, 2, 3, 4]。北京大學(xué)與字節(jié)跳動研究部門 ByteDance Research 聯(lián)合開發(fā)的計算框架 Forward Laplacian 創(chuàng)新地利用 Laplace 算子前向傳播計算，為 NNVMC 領(lǐng)域提供了十倍的加速，從而大幅降低計算成本，達成該領(lǐng)域多項 State of the Art，同時也助力該領(lǐng)域向更多的科學(xué)難題發(fā)起沖擊。該工作以《A computational framework for neural network-based variational Monte Carlo with Forward Laplacian》為題的論文已發(fā)表于國際頂級期刊《Nature Machine Intelligence》，相關(guān)代碼已開源。

• 論文鏈接：https://www.nature.com/articles/s42256-024-00794-x
• 代碼地址：
• https://github.com/bytedance/LapNet。
• https://github.com/YWolfeee/lapjax。

該項工作一提出即受到相關(guān)研究人員的密切關(guān)注，圍繞該工作已有多個開源項目實現(xiàn)，編程框架 JAX 也計劃將該項工作吸收其中。

該項工作由北京大學(xué)智能學(xué)院王立威課題組、物理學(xué)院陳基課題組聯(lián)合字節(jié)跳動研究部門 ByteDance Research 一同開發(fā)完成，作者中有多位北京大學(xué)博士生在 ByteDance Research 實習(xí)。

背景簡介

基于神經(jīng)網(wǎng)絡(luò)的量子變分蒙特卡洛方法（NNVMC）已成為量子化學(xué) - 從頭計算領(lǐng)域中一項前沿技術(shù)。它具備精度高、適用范圍廣等優(yōu)點。但它的阿克琉斯之踵在于過高的計算成本，這也限制了該方法在實際化學(xué)問題中的應(yīng)用。

作者提出了一套全新的計算框架 "Forward Laplacian"，利用 Laplace 算子的前向傳播，顯著提升了 NNVMC 方法的計算效率，為人工智能在微觀量子問題中的應(yīng)用打開了新的大門。

方法介紹

Forward Laplacian 框架

在 NNVMC 方法中，神經(jīng)網(wǎng)絡(luò)的目標(biāo)函數(shù)是微觀體系的能量，包括動能與勢能兩項。其中動能項涉及對神經(jīng)網(wǎng)絡(luò)的拉普拉斯算子的計算，這也是 NNVMC 中耗時最長的計算瓶頸。現(xiàn)有的自動微分框架在計算拉普拉斯算子時，需要先計算黑塞矩陣，再求得拉普拉斯項（即黑塞矩陣的跡）。而作者所提出的計算框架 "Forward Laplacian" 則通過一次前向傳播直接求得拉普拉斯項，避免了黑塞矩陣的計算，從而削減了整體計算的規(guī)模，實現(xiàn)了顯著加速。

LapNet 網(wǎng)絡(luò)

除了有效削減計算圖規(guī)模之外，F(xiàn)orward Laplacian 框架的另一大特點是能有效利用神經(jīng)網(wǎng)絡(luò)梯度計算中的稀疏性，提出神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu) LapNet。LapNet 通過增加神經(jīng)網(wǎng)絡(luò)中的稀疏性，在精度無損的同時，顯著提升了網(wǎng)絡(luò)計算的效率。

計算結(jié)果

絕對能量

作者首先就方法的效率及精度同當(dāng)前 NNVMC 領(lǐng)域有代表性的幾項工作進行了比較。從絕對能量的計算結(jié)果而言，作者提出的 LapNet 在 Forward Laplacian 框架下的效率高于參考工作數(shù)倍，精度上也與 SOTA 保持一致。此外，如果在相同計算資源（即相同 GPU hour）的情況下比較，LapNet 的計算結(jié)果可以顯著優(yōu)于之前的 SOTA。

加速標(biāo)度

為了更明確地研究作者所提出方法相比于之前 SOTA 的加速標(biāo)度，作者在不同大小的鏈?zhǔn)骄垡蚁w系上進行了測試，結(jié)果可以很明顯地看到 Forward Laplacian 工作帶來的 O (n) 加速。此處 n 為目標(biāo)分子中的電子數(shù)目。

相對能量

在物理、化學(xué)研究中，相對能量相較于絕對能量具有更明確的物理意義。作者也在一系列的體系上進行了測試，均取得了理想結(jié)果。

總結(jié)

為降低基于神經(jīng)網(wǎng)絡(luò)的量子變分蒙特卡洛方法（NNVMC）的使用門檻，北京大學(xué)與字節(jié)跳動研究部門 ByteDance Research 聯(lián)合開發(fā)了計算框架 Forward Laplacian，實現(xiàn)了十倍的加速。該工作已受到相關(guān)研究人員的廣泛關(guān)注，期望能夠推動 NNVMC 方法在更多科學(xué)問題中發(fā)揮重要作用。

參考文獻

[1] Han, J., Zhang, L., & Weinan, E. (2019). Solving many-electron Schr?dinger equation using deep neural networks. Journal of Computational Physics, 399, 108929.

[2] Hermann, J., Sch?tzle, Z., & Noé, F. (2020). Deep-neural-network solution of the electronic Schr?dinger equation. Nature Chemistry, 12 (10), 891-897.

[3] Pfau, D., Spencer, J. S., Matthews, A. G., & Foulkes, W. M. C. (2020). Ab initio solution of the many-electron Schr?dinger equation with deep neural networks. Physical Review Research, 2 (3), 033429.

[4] Li, X., Li, Z., & Chen, J. (2022). Ab initio calculation of real solids via neural network ansatz. Nature Communications, 13 (1), 7895.