WANet與WALoss：突破量子化學中的哈密頓量計算瓶頸

（本文閱讀時間：9分鐘）

編者按：量子化學中的哈密頓量計算一直是制約藥物篩選與新材料設計效率的瓶頸。近期，微軟研究院科學智慧中心、耶魯大學及加州大學戴維斯分校聯合研究提出的全新損失函式 WALoss 和深度學習框架 WANet，為這一難題提供了創新的解決方案。透過最佳化計算過程，WANet 顯著提高了哈密頓量預測的精度與效率，尤其在大分子系統中展示了強大的擴充套件能力。搭配 WALoss，WANet 則能夠準確預測分子能量，並有效捕捉分子的電子分佈，成為化學、材料科學等領域多維分析的可靠工具。WANet 與 WALoss 的提出有望推動量子化學的廣泛應用，助力高通量藥物篩選、新材料設計和基礎科學研究的創新進展。

近年來，密度泛函理論（Density Functional Theory, DFT）憑藉其高效性和準確性，已成為量子化學與材料科學研究中的核心工具，廣泛應用於預測分子電子結構、光譜特性及分子動力學模擬。DFT 的核心在於求解 Kohn-Sham 方程，透過構建 Kohn-Sham 哈密頓量獲取分子系統的關鍵量子態資訊，如 HOMO（最高佔據分子軌道）和 LUMO（最低未佔據分子軌道）能量、HOMO-LUMO 能隙以及總能量和光譜特性，這些資訊為研究分子振動、反應路徑和結構最佳化提供了重要基礎。

然而，傳統自洽場方法（Self-consistent field, SCF）在構建 Kohn-Sham 哈密頓量時具有 O(N³) 到 O(N⁴) 的計算複雜度（N 為電子數），在原子數量較多的大分子系統中，計算資源需求呈指數增長，導致應用受限。現有深度學習方法雖在小分子系統上有所突破，但在大分子系統上往往因資料稀缺和誤差積累導致非物理性預測結果，從而顯著降低了模型的可靠性與適用性。

對此，微軟研究院科學智慧中心、耶魯大學及加州大學戴維斯分校合作提出了全新的損失函式波函式對齊損失 WALoss（Wavefunction Alignement Loss）和現代化深度學習架構 WANet [1]，為提升 Kohn-Sham 哈密頓量的可擴充套件性與預測精度提供了全新的解決方案。這項研究在分子系統建模方面取得了顯著突破，助力破解大規模分子系統的計算難題。相關研究成果 Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems 已被全球機器學習領域頂會 ICLR 2025 評為 Spotlight 研究。

Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems

論文連結：

https://openreview.net/forum?id=twEvvkQqPS&noteId=rc63eJkois

傳統方法在大分子哈密頓量預測中存在嚴重偏差？

為了初步驗證現有網路在哈密頓量可擴充套件性上的不足，研究員們開發了一個全新的資料集 PubChemQH。該資料集基於 GPU 加速的密度泛函理論軟體 MADFT [2]，使用了 def2-tzvp 基組和 B3LYP 泛函，在128張 NVIDIA V100 顯示卡上進行了一個月的不間斷計算生成。PubChemQH 資料集包含50,000個分子樣本，樣本分子結構中的原子數量範圍從40到100，超過此前金標準資料集 QH9（僅包含不超過31個原子的分子）。這一大規模資料集的引入，為後續探索大分子系統中的可擴充套件性提供了基礎。

在 PubChemQH 資料集上，研究員們發現，儘管傳統深度學習方法在哈密頓量預測中表現出較低的平均誤差，但其預測的物理性質（如總能量、波函式、HOMO、LUMO 等）偏差極大。其中，總能量的預測誤差高達約10,000 kcal/mol，遠遠無法達到化學精度（通常要求誤差在 1 kcal/mol 以內），嚴重影響了模型的實際可用性。

突破哈密頓量計算瓶頸

針對預測偏差大的問題，研究員們提出了全新的損失函式波函式對齊損失 WALoss 和現代化深度學習架構 WANet。WALoss 旨在透過對齊預測波函式與真實波函式的基，來提升模型預測的物理準確性。由於傳統方法大多采用均方誤差（MSE）或平均絕對誤差（MAE）作為損失函式，這些損失函式僅對哈密頓矩陣上的每一項進行單獨比較，所以無法有效捕捉哈密頓量在整體結構上的偏差，尤其在處理大分子時往往會失效。WALoss 可透過基變換對波函式進行對齊，使得預測的哈密頓量更加符合物理真實值，從而減少預測誤差並顯著提升 SCF 計算的收斂速度。

專門用於 Kohn-Sham 哈密頓量預測的現代化網路架構 WANet 則進一步提升了網路的計算效率。WANet 的構建靈感來源於等變圖神經網路 SE(3)-Equivariant Networks，該網路能夠自動捕捉分子的旋轉、平移及對稱性等物理特徵，有效提高了模型對分子結構的理解能力。

在 WANet 的網路架構中，節點卷積層（Node Convolution Layer）承擔著捕捉區域性原子間複雜相互作用的任務。透過對區域性原子鄰域資訊的圖卷積操作，該層能夠生成節點的高階不可約表示，為下游的哈密頓量預測構建模組提供支援。具體地，WANet 將之前網路中常使用的 SO(3) 卷積改為了基於 SO(2) 的卷積神經網路，保證了網路等變的同時，大幅提升了計算效率。最後，哈密頓量預測模組則利用 Clebsch-Gordan 張量積技術，將高階不可約表示進一步轉化為完整的哈密頓量矩陣，並分別構建對角塊和非對角塊。

圖1：WANet 研究流程

WANet 和 WALoss 實驗驗證

實驗結果表明，WANet 和 WALoss 在 PubChemQH 資料集上取得了極大的效能提升。在 PubChemQH 資料集上（如表1），WANet 搭配 WALoss 後，系統能量預測誤差從 QHNet 的65721.028 kcal/mol 降至47.193 kcal/mol，同時 SCF 迭代次數減少至82%，遠低於傳統方法的371%。

表1：WANet 和 WALoss 在 PubChemQH 上的實驗結果

計算效率是深度學習在量子化學應用中的重要考量。WANet 在推理速度和 GPU 資源利用方面展現出顯著優勢。相比傳統 DFT 方法，WANet 不僅減少了 SCF 迭代次數，還大大降低了總計算時間。GPU 記憶體佔用也從 QHNet 的26.49 GB 降至15.86 GB，使得 WANet 更加適合高通量計算場景。

圖2：WANet 和 WALoss 的計算效率

除了能量預測，WANet 在偶極矩（dipole moment）和電子空間範圍（Electronic Spatial Extent, ESE）的預測中也表現優異。實驗資料顯示，WANet 搭配 WALoss 後在 Dipole Moment 預測中的誤差僅為3.3928 D，顯著優於其他基於深度學習的迴歸模型。同時，在 ESE 預測中，WANet 的誤差為0.0076 a.u.，遠低於傳統方法。這一結果證明 WANet 不僅在能量預測方面具有優勢，還能有效捕捉分子的空間電子分佈，為分子性質的多維度分析提供了可靠工具。

表2：偶極矩和電子空間範圍預測

WANet 還具有出色的擴充套件能力。在碳鏈延伸分子（CxH2x+2）的測試中，當碳原子數量超過 PubChemQH 資料集的範圍（超過100個原子）時，WANet 依然保持了高精度預測，尤其在 LUMO 能量的預測中顯著優於傳統方法。這一結果表明，WANet 能夠有效泛化到更大規模的分子系統。

未來應用潛力

WANet 和 WALoss 的提出為密度泛函理論的進一步發展注入了技術活力，也為高通量藥物篩選、新材料設計及基礎科學研究開闢了全新路徑。隨著更高質量資料集的引入和模型架構的不斷最佳化，基於深度學習的哈密頓量預測方法有望在未來實現更廣泛的應用，推動量子化學從理論走向實踐，助力基礎科學與工業技術的進步。

相關連結：

[1] https://openreview.net/forum?id=twEvvkQqPS

[2] Ju F, Wei X, Huang L, et al. Acceleration without disruption: DFT software as a service[J]. Journal of Chemical Theory and Computation, 2024, 20(24): 10838-10851.

你也許還想看：