基于神經網絡的偏微分方程求解器新突破：北大&字節研究成果入選Nature子刊
近年來，基于神經網絡的偏微分方程求解器在各領域均得到了廣泛關注。其中，量子變分蒙特卡洛方法（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 中耗時最長的計算瓶頸?，F有的自動微分框架在計算拉普拉斯算子時，需要先計算黑塞矩陣，再求得拉普拉斯項（即黑塞矩陣的跡）。而作者所提出的計算框架 "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.
