低成本、準確、穩健，各類分子通用，上海人工智慧實驗室開發MD模擬AI新方法

ScienceAI發表於2024-04-12
人工智慧AI

圖片

編輯 | 綠蘿

機器學習原子間勢（MLIP）因其兼顧高精度和高效率的優勢，在材料、化學、生物學等領域的大尺度原子模擬研究中引起了廣泛關注。

然而，高效能 MLIP 依賴於大量標記資料，透過從頭計算獲取這些資料的成本很高。

近日，上海人工智慧實驗室、復旦大學和清華大學的研究團隊，提出了一種 MLIP 的幾何學習框架 GPIP，利用未標記的構型來提高 MLIP 的效能。

研究表明，GPIP 只需少量的計算成本即可顯著提高 MLIP 的準確性和泛化性，並且與不同的不變或等變圖神經網路架構相容。該方法增強了 MLIP，並推進了分子系統的模擬。

相關研究以「Geometry-enhanced pretraining on interatomic potentials」為題，於 2024 年 4 月 5 日釋出在《Nature Machine Intelligence》上。

圖片

論文連結：https://www.nature.com/articles/s42256-024-00818-6

分子動力學 (MD) 在物理、化學、生物和材料科學等諸多領域中提供了原子見解。MD 模擬的準確性和效率取決於原子間勢的選擇，原子間勢是描述分子系統中原子勢能的數學函式。

機器學習原子間勢（MLIP）能夠藉助機器學習模型擬合從頭算能量和力，從而準確高效地實現從頭算。其中，圖神經網路被廣泛使用，研究者還開發了等變神經網路來保持分子系統的等變性，以提高原子間勢預測的效能。

儘管 MLIP 已被廣泛研究，但由於從頭算的成本高昂，訓練資料的缺乏，其效能和可轉移性受到限制。

GPIP：機器學習原子間勢的幾何學習框架

在此，研究團隊提出了一種 MLIP 的幾何學習框架：原子間勢幾何增強預訓練（Geometry-enhanced Pretraining on Interatomic Potentials，GPIP） ，該框架由幾何結構生成和幾何增強預訓練兩部分組成。

圖片

圖示：GPIP 的總體架構。（來源：論文）

研究人員提出了一個自監督學習（SSL）框架，首先對廉價且易於獲取的經典 MD (CMD) 資料進行自監督預訓練，然後對從頭算 MD (AIMD) 生成的標記資料進行微調以學習任務相關資訊。在此框架下，可以以有限的成本輕鬆生成特定分子系統的大規模預訓練資料集，增強預訓練能力，並避免預訓練資料系統與測試資料系統之間的域差距。

在此基礎上，研究人員還提出了一種幾何增強的 SSL 方法，該方法涉及三個互補的任務，包括使用噪聲座標恢復遮蔽原子、使用遮蔽原子預測噪聲以及使用 3D 網路進行對比學習，以便更好地從 CMD 資料中提取拓撲和 3D 幾何資訊。

該研究的主要貢獻如下：

（1）提出了一個用於 MLIP 的幾何學習框架：GPIP，由兩個部分組成：幾何結構生成和幾何增強預訓練。

（2）研究證明，CMD 模擬產生的未標記幾何結構有利於原子間勢預測，且計算成本非常低。

（3）設計了一種針對分子結構資料的幾何增強自監督預訓練方法。

（4）研究人員使用涵蓋各種分子系統的 MLIP 和資料集來評估其方法。該方法在所有實驗中都表現出一致的穩健性。

GPIP 的有效性歸因於自監督任務能夠從具有經驗力場的 CMD 模擬生成的資料中有效提取結構資訊的能力。儘管不如從頭計算方法準確，但 CMD 模擬可以提供物理上合理的原子結構，因為經驗力場基於物理公式。

圖片

圖示：將預訓練和微調資料投影到 SchNet-GPIP 模型的嵌入上。（來源：論文）

此外，當從不同溫度和壓力下的長時間 CMD 模擬中取樣時，可以從大空間取樣豐富的結構資料，覆蓋下游任務中的微調資料集。這種覆蓋範圍可以透過均勻流形近似和投影降維方法視覺化。

涵蓋各種分子系統

為了證明其方法的有效性，研究人員用 GPIP 預訓練 MLIP，並將其應用於各種具有挑戰性的基準以進行原子間勢預測。

表 1：有機小分子 MD17 資料集的結果。（來源：論文）

圖片

MD17 資料集由 AIMD 模擬生成的有機小分子組成。基於 GPIP 的模型與基線模型的力預測平均絕對誤差（MAE）比較如表 1 所示。基於 GPIP 的 SchNet、DimeNet、SphereNet 和 GemNet-T 的效能均優於相應的基線模型 。

圖片

圖示：在 ISO17 資料集上的結果。（來源：論文）

ISO17 資料集由 AIMD 模擬生成的 C7O2H10 異構體的 MD 軌跡組成。研究人員首先使用 ISO17 資料集中的 20,000 個構型的小資料集和 400,000 個構型的大資料集進行微調。結果顯示，儘管微調資料集非常大，但基於 GPIP 的模型在這兩項任務上都顯著優於非預訓練的模型。

表 2：Water 資料集的結果。（來源：論文）

圖片

由於結構的複雜性，為帶有周期性邊界條件（PBCs）的擴充套件系統訓練 MLIP 比單分子系統更困難。研究人員將 GPIP 用於含有 PBC 的液態水研究。基於 GPIP 的模型不僅提高了力預測的準確性，而且還提高了 MD 模擬中的其他關鍵指標。

表 3：電解質資料集的結果。（來源：論文）

圖片

為了更好地評估 MLIP 的效能，還開發了一個電解質溶液資料集，其中包含更多型別的元素和構型。

值得強調的是，與獲得訓練集準確標籤的從頭計算相比，GPIP 的計算複雜度可以忽略不計。雖然增加訓練資料的數量也可以提高準確性，但由於從頭計算，計算成本可能比 GPIP 高數百倍。

總之，研究人員提出了一種稱為 GPIP 的 MLIP 幾何學習框架。實驗結果表明，基於 GPIP 的模型始終顯著優於相應的基線模型，表現出顯著的有效性和穩健性。

該研究非常適用於使用 MLIP 進行 MD 模擬，特別是對於難以獲得足夠訓練資料的複雜分子系統。

注：封面來自網路

相關文章