近年來，基於神經網路的偏微分方程求解器在各領域均得到了廣泛關注。其中，量子變分蒙特卡洛方法（NNVMC）在量子化學領域異軍突起，對於一系列問題的解決展現出超越傳統方法的精確度 [1, 2, 3, 4]。北京大學與位元組跳動研究部門 ByteDance Research 聯合開發的計算框架 Forward Laplacian 創新地利用 Laplace 運算元前向傳播計算，為 NNVMC 領域提供了十倍的加速，從而大幅降低計算成本，達成該領域多項 State of the Art，同時也助力該領域向更多的科學難題發起衝擊。該工作以《A computational framework for neural network-based variational Monte Carlo with Forward Laplacian》為題的論文已發表於國際頂級期刊《Nature Machine Intelligence》，相關程式碼已開源。
論文連結：https://www.nature.com/articles/s42256-024-00794-x
程式碼地址：
https://github.com/bytedance/LapNet
https://github.com/YWolfeee/lapjax
該項工作一提出即受到相關研究人員的密切關注，圍繞該工作已有多個開源專案實現，程式設計框架 JAX 也計劃將該項工作吸收其中。
該項工作由北京大學智慧學院王立威課題組、物理學院陳基課題組聯合位元組跳動研究部門 ByteDance Research 一同開發完成，作者中有多位北京大學博士生在 ByteDance Research 實習。
背景簡介
基於神經網路的量子變分蒙特卡洛方法（NNVMC）已成為量子化學 - 從頭計算領域中一項前沿技術。它具備精度高、適用範圍廣等優點。但它的阿克琉斯之踵在於過高的計算成本，這也限制了該方法在實際化學問題中的應用。
作者提出了一套全新的計算框架 "Forward Laplacian"，利用 Laplace 運算元的前向傳播，顯著提升了 NNVMC 方法的計算效率，為人工智慧在微觀量子問題中的應用開啟了新的大門。
方法介紹
Forward Laplacian 框架
在 NNVMC 方法中，神經網路的目標函式是微觀體系的能量，包括動能與勢能兩項。其中動能項涉及對神經網路的拉普拉斯運算元的計算，這也是 NNVMC 中耗時最長的計算瓶頸。現有的自動微分框架在計算拉普拉斯運算元時，需要先計算黑塞矩陣，再求得拉普拉斯項（即黑塞矩陣的跡）。而作者所提出的計算框架 "Forward Laplacian" 則透過一次前向傳播直接求得拉普拉斯項，避免了黑塞矩陣的計算，從而削減了整體計算的規模，實現了顯著加速。
LapNet 網路
除了有效削減計算圖規模之外，Forward Laplacian 框架的另一大特點是能有效利用神經網路梯度計算中的稀疏性，提出神經網路結構 LapNet。LapNet 透過增加神經網路中的稀疏性，在精度無損的同時，顯著提升了網路計算的效率。
計算結果
絕對能量
作者首先就方法的效率及精度同當前 NNVMC 領域有代表性的幾項工作進行了比較。從絕對能量的計算結果而言，作者提出的 LapNet 在 Forward Laplacian 框架下的效率高於參考工作數倍，精度上也與 SOTA 保持一致。此外，如果在相同計算資源（即相同 GPU hour）的情況下比較，LapNet 的計算結果可以顯著優於之前的 SOTA。
加速標度
為了更明確地研究作者所提出方法相比於之前 SOTA 的加速標度，作者在不同大小的鏈式聚乙烯體系上進行了測試，結果可以很明顯地看到 Forward Laplacian 工作帶來的 O (n) 加速。此處 n 為目標分子中的電子數目。
相對能量
在物理、化學研究中，相對能量相較於絕對能量具有更明確的物理意義。作者也在一系列的體系上進行了測試，均取得了理想結果。
總結
為降低基於神經網路的量子變分蒙特卡洛方法（NNVMC）的使用門檻，北京大學與位元組跳動研究部門 ByteDance Research 聯合開發了計算框架 Forward Laplacian，實現了十倍的加速。該工作已受到相關研究人員的廣泛關注，期望能夠推動 NNVMC 方法在更多科學問題中發揮重要作用。
參考文獻
[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.