近年來，基于神經(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ā)的計(jì)算框架 Forward Laplacian 創(chuàng)新地利用 Laplace 算子前向傳播計(jì)算，為 NNVMC 領(lǐng)域提供了十倍的加速，從而大幅降低計(jì)算成本，達(dá)成該領(lǐng)域多項(xiàng) State of the Art，同時(shí)也助力該領(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。

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

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

背景簡介

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

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

方法介紹

Forward Laplacian 框架

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

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

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

計(jì)算結(jié)果

絕對能量

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

加速標(biāo)度

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

相對能量

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

總結(jié)

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

參考文獻(xiàn)

[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.