編輯 | ScienceAI
編者按：為了使電子結構方法突破當前廣泛應用的密度泛函理論（KSDFT）所能求解的分子體系規模，微軟研究院科學智慧中心的研究員們基於人工智慧技術和無軌道密度泛函理論（OFDFT）開發了一種新的電子結構計算框架M-OFDFT。這一框架不僅保持了與KSDFT相當的計算精度，而且在計算效率上實現了顯著提升，並展現了優異的外推效能，為分子科學研究中諸多計算方法的基礎——電子結構方法開闢了新的思路。相關研究成果已在國際知名學術期刊《自然-計算科學》（Nature Computational Science）上發表。
然而，各種電子結構方法的求解精度和計算效率往往無法兼得。當前，取得相對合理的「精度-效率」權衡而被廣泛應用的方法是Kohn-Sham形式的密度泛函理論（Kohn-Sham density functional theory, KSDFT）。但KSDFT具有較高的計算複雜度，不能很好地滿足日益增長的求解大規模分子體系的需求。
為此，微軟研究院科學智慧中心的研究員們提出了一種基於深度學習和無軌道密度泛函理論（OFDFT）的電子結構計算框架M-OFDFT，其不僅顯著超越了KSDFT的計算效率，還能保有其求解精度。這一成果展示了人工智慧在提升電子結構計算中「精度-效率」權衡方面的卓越能力，並將助力加速相關業界問題的研究與開發。
M-OFDFT的相關研究成果以「Ovecoming the Barrier of Orbital-Free Density Functional Theory for Molecular Systems Using Deep Learning」為題，於 2024 年 3 月 11 日發表在國際知名學術期刊《自然-計算科學》（Nature Computational Science）上。
人工智慧給電子結構方法帶來新機會
電子結構方法是求解分子體系各種物理化學性質的基礎工具。由於多電子體系本身具有一定的求解難度，所以高精度電子結構方法因其較高的計算代價很難應用到工業界所關注的分子體系中，而可計算較大分子的方法則會因引入一些近似而損失精度。目前KSDFT因其相對合適的精度與效率權衡得到了廣泛應用。
不過，近期人工智慧技術的喜人進展也為其他電子結構計算框架帶來了新的機會。為了使電子結構方法突破KSDFT所能求解的分子體系規模，微軟研究院的研究員們利用人工智慧技術，開發了M-OFDFT，該方法比KSDFT效率更高，同時又能保有其精度。基於OFDFT的開發，讓M-OFDFT成為了一種比KSDFT理論複雜度更低的電子結構計算框架，因為它只需最佳化電子密度函式這一個函式來求解電子狀態即可，KSDFT則需要最佳化與電子數相同的多個函式。
不過，OFDFT面臨著一個巨大的挑戰——需要電子動能關於密度函式的泛函，但它的形式未知，並且難以構造適用於分子體系的高精度近似。
針對這一難題，M-OFDFT使用一個深度學習模型來近似動能泛函。藉助深度學習模型的強大擬合能力，M-OFDFT可實現比基於近似物理模型設計的經典動能泛函更高的準確度。對於一個待求解的分子體系結構，M-OFDFT會使用動能泛函模型以及其他可直接計算的能量項構造出一個電子密度的最佳化目標，然後透過最佳化過程求解最優（基態）電子密度（圖1），進而可計算能量、力、電荷分佈等分子屬性。
圖1：對於待求解的分子體系結構，M-OFDFT透過最小化電子能量來求解電子密度（以其向量化係數表示），其中難以近似的動能部分由深度學習模型來近似。
M-OFDFT實現兼具精度與效率的電子結構方法
研究員們對M-OFDFT進行了一系列的實驗驗證。首先考察的是M-OFDFT在常見小分子體系上的求解精度。
結果顯示，M-OFDFT在乙醇分子構象以及QM9資料集的分子上可以達到與KSDFT相當的精度（能量達到化學精度1 kcal/mol）。相較於經典OFDFT方法，精度提高了兩個數量級（圖2-a）。M-OFDFT解得的電子密度也與KSDFT的結果重合（圖2-b），特別是得到了電子殼層結構，而經典OFDFT的結果則有明顯偏差。由M-OFDFT解得的乙醇構象空間上的勢能面（每個點都是透過密度最佳化得到的，並不是直接預測）也與KSDFT的結果一致（圖2-c）。
之後，研究員們又驗證了M-OFDFT不僅保有KSDFT級別的精度，其更低的理論計算複雜度還使其在效率上也超越了KSDFT。在實際計算中M-OFDFT取得了的複雜度（圖3），比KSDFT 的實際複雜度低了一階，且其所需絕對時間也明顯少於KSDFT。在兩個更大的蛋白質體系上（包含2676和2750個電子），M-OFDFT實現了25.6倍和27.4倍的加速。
M-OFDFT具有更強的泛化能力
深度學習模型在科學任務中的應用面臨一大挑戰是，在具有與訓練資料不同特點的資料上的泛化問題。但採用了OFDFT框架後，動能泛函模型遇到的泛化問題就會減輕，從而使 M-OFDFT可以在比訓練集分子規模更大的體系上展現出良好的外推能力。
實驗結果表明，M-OFDFT的能量預測誤差顯著低於基於深度學習的端到端能量預測模型（圖4-a）。此外，研究員們還利用在多肽片段上訓練的M-OFDFT模型求解完整蛋白結構，並取得了超越端到端模型和經典OFDFT的泛化效能（圖4-c）。不僅如此，相較端到端模型，M-OFDFT還可以用更少的大分子體系訓練資料取得更好的泛化表現（圖4-b與圖4-d）。
M-OFDFT的工作原理
「神龍見首又見尾」：高效捕獲非局域效應的動能泛函模型
動能密度泛函具有明顯的非局域效應，而用經典的基於格點（grid）的方式表徵電子密度則會帶來高昂的非局域計算代價。為此，M-OFDFT將電子密度在一組原子基組函式上展開，並使用展開係數作為電子密度表徵。由於基函式疊加的形狀與電子分佈接近，所以其數量可遠小於格點數，使得非局域計算代價大大降低，並有助於刻畫電子密度中的殼層結構。
M-OFDFT將每個原子上的電子密度係數和型別與座標作為節點特徵,並基於Graphormer模型 [1]預測電子動能（圖5），其自注意力機制顯式刻畫了荷載在每兩個原子上的電子密度特徵之間的相互作用，從而可捕捉非局域性質。此外，為了保證動能的旋轉不變性，M-OFDFT使用了以各個原子為中心、基於其相鄰原子的區域性座標系，將電子密度係數轉換為旋轉不變的特徵。
「橫看成嶺側成峰，遠近高低各不同」：高效學習電子能量曲面的訓練策略
與傳統機器學習任務不同，動能泛函模型是被當作其輸入變數的最佳化目標使用的，而非用於在一些單點上做預測，這對模型的學習提出了更高的要求：模型必須捕捉到每個分子結構上電子能量曲面的輪廓。
為此，研究員們深入分析了用來生成資料的電子結構方法，發現它其實可以為每個分子結構生成多個資料點，而且還能提供梯度標註，從而讓模型可以擁有更豐富的曲面輪廓特徵。然而梯度的巨大範圍也使神經網路難以最佳化。對此，研究員們還提出了一系列增強模組，讓模型能夠更容易地表達巨大的梯度。
開啟未來電子結構方法的新篇章
M-OFDFT成功突破了無軌道密度泛函框架在分子體系中的瓶頸，將其求解精度提升到了常用的KSDFT的水平，同時保有了其更低的計算代價，推進了電子結構方法在「精度-效率」方面的權衡，為分子科學研究提供了一種更有潛力的研究工具。
儘管M-OFDFT已經在某些分子體系上展現了出色的泛化效能，但在更大的分子體系上實現長時間且穩定的高精度模擬仍是一個巨大的挑戰。微軟研究院期待M-OFDFT可以沿著這一方向激發更多研究與創新，並在未來和其他方法一起為電子結構計算帶來更多突破性的成果和影響。
相關文章：
[1] Do Transformers really perform badly for graph representation? Advances in Neural Information Processing Systems 34 (NeurIPS 2021).（https://proceedings.neurips.cc/paper/2021/hash/f1c1592588411002af340cbaedd6fc33-Abstract.html）
注：封面來自網路。