作者 | 楊喜喜
審稿 | 姜晶
今天給大家介紹來自牛津大學克拉倫登實驗室和俄羅斯量子力學中心的A.I.Lvovsky團隊與來自InstaDeeP公司和牛津大學克拉倫登實驗室的Thomas D. Barrett合作發表在《Nature Machine Intelligence》的論文《Autoregressive neural-network wavefunctions for ab initio quantum chemistry》。該論文提出使用神經網絡參數化波函數，引入高效且可擴充的采樣，在嵌入反映分子系統結構的實體先驗的同時，保留網絡的表達能力。
1
背景介紹
量子化學（QC）第一性原理計算需要求解重原子核和軌道電子互相作用的多體薛定谔方程。原則上，求解薛定谔方程可以完整的描述分子的化學性質。然而，超乎尋常的系統複雜性，使得多體電子結構計算是一個NP-hard問題。為了解決這個問題，科學家們緻力于将數值方法應用于量子化學系統。但電子結構複雜的非局部互動特性和高采樣難度，給使用神經網絡量子态計算帶了極大的困難。
神經網絡量子态計算旨在計算具有Ne個電子的分子系統的基态。使用全構型互相作用(Full configuration interaction, FCI)來計算分子系統的基态，需要計算電子占據軌道的反對稱張量積（Slater行列式）。當系統的大小增長時，Slater行列式的可能數量呈指數增長。這樣，使用FCI計算Slater行列式迅速變得難以處理。量子化學第一性原理通常對參考狀态(Hartee Fock狀态，即單個能量最低的Slater行列式)進行系統修正。一般有三種方法，1)構型互動方法，将電子軌道的可能活動空間限制為構型字元。且這些構型的字元有别與不超過一定激勵的參考字元；2)耦合聚類(coupled-cluster, CC)，使用一定順序的非線性激勵算子組合來表示任意激勵的Slater行列式。但耦合聚類方法不能參數化任意量子态的疊加。這兩種方法都犧牲了表達能力以降低系統的複雜性，但他們都是基于典型結構的波函數設計的，通常提供了良好的表達能力。然而，即使是最精确的CC方法在系統存在強靜态相關時(即當對單個Slater行列式的系統矯正無法充分描述波函數時)仍然可能失敗；3)對波函數進行緊湊的參數化，例如Jastrow-Slater ansatz或矩陣積态，然後利用随機或非随機方法對其進行優化，以找到基态。
2017年，Carleo和Troyer證明，基于受限玻爾茲曼機(RBM)的神經網絡可以參數化一個多體波函數，并捕獲指數大編碼希爾伯特空間内的非瑣碎相關性。随後，神經網絡量子态(NNQSs)已被證明是一些問題的可行變分 ansatzes(拟設)，如尋找有互相作用自旋系統的基态、量子态層析成像和量子計算的經典模拟。大多數NNQSs的發展一直在凝聚态實體(CMP)的背景下，近期研究表明費密子電子結構問題可以被映射一個互相作用的旋轉系統的等效優化問題，為将NNQSs應用于量子化學打開了大門。然而，RBM模型的采樣程式（如馬爾可夫鍊，蒙特卡羅方法近似的玻爾茲曼分布編碼）低效。此外，RBM是一個黑匣子，不能反映分子系統的實體先驗。是以，仍需尋找克服這些缺點的神經網絡架構。
在這項工作中，作者提出神經網絡模拟分子二次量化的波函數，以解決上文提到的挑戰。通過考慮波函數的自回歸分解，本文示範了一種高效的采樣算法，該算法固有地并行化，并根據采樣的唯一構型字元串數量縮放。使用前饋結構，使得模型可以使用标準的反向傳播技術來訓練網絡，并将重要的實體先驗(即守恒量和不變性)嵌入到波函數中，而不犧牲其表達能力。最終，這使NAQS能夠在比傳統NNQS至少大25倍的系統上接近FCI效果。
2
模型結構
如圖1所示，模型的自回歸網絡如圖1(a)所示，使用Subnetwork(圖1(b))對空間布局字元
的四種可能的立體基陣的振幅進行對數規範化。然後進行采樣,作為下一階段網絡的輸入資料。ARN網絡中第i個子網絡接受部配置設定置字元串
作為輸入。輸出四種可能構型振幅的正則化
，該網絡可用于兩種不同的模式，這裡稱之為推理和采樣。
圖1.ARN網絡結構
模型的推理和采樣的架構如圖2所示。如圖2(a)所示在推理階段，模型評估給定構型字元串xk的對數振幅，以及該構型字元串對應于網絡的單個前饋傳遞。如圖2(b)所示，采樣的任務是根據底層的分布
生成xk。該過程是為了輸出唯一的構型字元串而定制的，其中每個字元串都與一個數字nk相關聯，nk表示如果使用标準過程(包括對配置進行N次獨立采樣)，它将發生多少次。
圖2.自回歸量子态（NAQS）模型架構
3
結果分析
如表1所示，NAQS在所有分子上都實作了強勁的預測能力。作者使用Hartree Fock能量(HF)、單激勵和雙激勵下的構型互相作用計算(CISD)，雙激發耦合簇方法(CCSD)和三激發耦合簇方法(CCSD(T))作為近似量子化學計算的baseline。如表1所示，與baseline相比，NAQS在所有分子上都表現出優異的性能。在分子(LiCl)具有高達20個電子和28個自旋軌道的情況下，即使最複雜的baseline已經失效，NAQS依然能獲得接近或符合FCI結果。事實上，即使在具有40M個實體上有效的基行列式的分子(Li2O)的FCI計算上，NAQS仍然優于CCSD和CISD。一個明顯的例外是，CCSD(T)預測的Li2O能量較低;這個結果也低于FCI極限(是以是非實體的)。相反NARN ansatz保證其輸出狀态的實體有效性。
表1.使用不同baseline和NAQS的變體獲得分子的基态能(機關Ha)
圖3進一步說明了NAQS的優越性，在圖3中，即使在耦合激發簇(coupled-cluster，CC)基線失效的區域，NAQS在N2勢能面計算上也保持了與FCI-like一緻的準确性。
圖3.利用NAQS方法和傳統QC方法獲得的雙原子氮分子的能量比較
表2.NAQS與RBM的性能比較
為了直接比較NAQS與RBM的性能，作者将NAQS應用于RBM中給出的分子幾何形狀，這些結果總結在表2中。
如表2所示，将NAQS的結果與Choo等人使用受限玻爾茲曼機(RBM)在二次量化分析中将神經網絡量子态(NNQS)應用于QC的最好的(也是唯一的)方法進行比較，具有啟發性。RBM的神經網絡量子态超過CC方法的最大系統是C2 (44.1k個行列式)，在具有約1.6M個行列式的系統上，對較大分子的限制采樣性能優于CISD。NAQS的ansatz可以擴充到比H2大一個數量級以上的系統而不損失性能，并且在每個比H2大的分子上都優于RBM(NAQS在分子H2上取得了和RBM相當的結果)。
作者還觀察到大分子優化空間采樣不足的問題，具體來說，在遮蔽非實體行列和不遮蔽非實體行列式的情況下，對于LiCl和Li2O，NAQS通常在訓練期間至少一次分别隻采樣～28%(20%)和～5%(1%)的實體行列式，仍然學習到精确波函數所需的構型資訊。這表明，神經網絡可以很好的生成新的構型,并且提供了一種有效地探索和隔離整個優化空間的重要元件的方法。
經過一段初期的搜尋，自回歸神經網絡(autoregressive neural network, ARN)通常分離出少數主導模型表現的決定性因素。之後的優化過程會微調這些決定因素的振幅。例如，即使沒有用掩碼将ARN限制為僅實體可行性的決定因素，Li2O的最後一步也可以生成1012個樣本，僅分布在13,761個唯一決定因素上。如果取出由這些行列式張成的子空間，并簡單地對角化哈密頓量，可以得到87.8911Ha的能量，這甚至比ARN發現的能量(87.8909Ha)更低。這表明使用混合方法，将規範數值方法應用于網絡發現的構型的子空間，可以實作更低的能量預測。
燒蝕實驗：
為了檢查将實體先驗編碼對NAQS的影響，表1還包括了兩個消融的結果。無掩碼的NAQS可以為
個可能配置中的任何一個配置設定非零機率，進而對其進行采樣，無論它們是否具有實體可行性。“No spin sym”表示從波函數中去除自旋翻轉不變性。一個例外是具有開殼層的分子(CH2和O2)，NAQS無法表示他們的自旋翻轉不變性。然而，即使沒有這個限制，NAQS也能表現出堪比FCI-level的性能。
圖4. NAQS模型變體在優化過程中獲得的效能(标準)，以及将優化空間限制為實體可行行列(No mask)和自旋翻轉對稱性的相關消融
雖然移除這兩種實體先驗中的任何一種都會導緻性能的輕微下降，但影響是微小的。然而，當我們考慮多次初始化的性能時，差別是顯而易見的。圖4顯示了在執行多次優化的6個大分子上，每次運作獲得的能量，以及能量的平均值。當實體知識被編碼到系統中時，平均能量得到了改善。此外，在少數情況下，最好的結果是在消融實驗中得到的(例如，LiF上沒有掩碼)，很明顯，這是一個偶然的結果，而不是系統的效果。一般來說，這些燒蝕的影響隻在較大的系統上可見，如圖4所示。這似乎是合理的，因為随着希爾伯特空間次元的增加，可能的配置字元串集呈指數增長，是以在已知的先驗限制下，減小其大小變得更加關鍵。
4
總結
該論文提出了一種基于量子化學第一性原理的ARN ansatz，它在多達30個自旋軌道的分子系統獲得了接近FCI能量。即使在正則近似方法失敗的情況下，性能依然強勁。這一成功歸功于将實體先驗嵌入網絡，以及高效的采樣程式。該程式可以根據采樣的唯一構型數量而不是整個批處理大小進行縮放。最終，該系統擴充方法遠遠超出了之前基于機器學習二次量化QC方法。
資料和代碼
https://github.com/tomdbar/naqs-for-quantum-chemistry
參考資料
Barrett, T.D., Malyshev, A. & Lvovsky, A.I. Autoregressive neural-network wavefunctions for ab initio quantum chemistry. Nat Mach Intell 4, 351–358 (2022).
https://doi.org/10.1038/s42256-022-00461-z