LSR-MP：突破分子模擬瓶頸，探索知識引導的AI正規化

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

編者按：分子動力學模擬在新藥開發、材料設計等領域發揮著重要作用。近年來機器學習技術的不斷發展，使得其對分子間相互作用的刻畫也更加精確，但卻面臨著隨分子體系擴大，計算效率降低和長程資訊丟失的難題。在此背景下，微軟研究院科學智慧中心同耶魯大學、西安交通大學提出了一種名為 LSR-MP 的新型分子動力學機器學習框架。該框架結合了物理洞見和幾何深度學習，透過在原子/分子片段上分別建模短程和長程效應，為大規模分子體系的高精度、高效模擬開闢了新的途徑。

分子動力學模擬是研究分子體系的“顯微鏡”，在新藥開發、材料設計等領域發揮著重要作用。但傳統的模擬方法面臨計算成本高、難以處理大分子等瓶頸。近年來，人工智慧技術的發展為突破這些瓶頸帶來了新的希望。

近日，微軟研究院科學智慧中心（Microsoft Research AI for Science）攜手耶魯大學、西安交通大學提出了一種名為 LSR-MP 的新型分子動力學機器學習框架。該框架巧妙融合了量子化學領域的分割思想與幾何深度學習技術，在精準捕捉分子間相互作用的同時，大幅提升了計算效率。相關論文已被國際人工智慧頂級學術會議之一的 ICLR 2024 接收。

論文連結：

https://openreview.net/pdf?id=rvDQtdMnOl

探索多尺度相互作用建模的可擴充套件性

分子動力學模擬是透過求解物理方程，預測分子體系的時間演化過程，來獲得其性質與功能的資訊。傳統的從頭算方法，如密度泛函理論(DFT)，雖然精度較高，但計算成本隨分子尺寸呈O(N³)到O(N⁴)之間的多項式級增長，難以應用於蛋白質等大分子體系。

近年來，研究人員嘗試利用機器學習來“學習”分子間的相互作用，從而加速模擬。這其中，圖神經網路（GNNs）憑藉強大的圖結構資料處理能力嶄露頭角。特別是等變圖神經網路（EGNNs）透過編碼分子體系的物理對稱性，在小分子資料集上展現出了優越的效能。然而，當把它們用於數百個原子的生物大分子時，研究人員發現長程相互作用(如靜電作用、範德華力)的資訊容易丟失，影響模型效能，而且簡單增大感受野或加深網路，會導致計算效率降低、過擬合、過度擠壓等問題發生。

如何在建模分子間多尺度相互作用的同時，保持模型的可擴充套件性，是一個亟待解決的難題。

在過去幾年，基於分子分割的量子化學方法和機器學習方法為大型分子體系的模擬帶來了新的曙光。基於分割的方法，如分治法（Divide-and-Conquer，DC）和基於共軛帽的分子分割法（Molecular Fractionation with Conjugated Caps，MFCC）可以將體系切割為易處理的小片段，在一定程度上降低了計算複雜度。然而，這類方法在描述片段間的多體相互作用時還不夠精確，限制了它們的進一步應用。

因此，一種既能像基於分割的方法那樣處理大體系，又能像等變神經網路那樣精準學習原子間複雜相互作用的新的機器學習框架亟需開發。LSR-MP框架由此誕生。

LSR-MP：融合物理洞見的分子動力學學習正規化

LSR-MP 框架結合了物理洞見和幾何深度學習，可以看作是對現有 EGNNs 的擴充套件。如圖1所示，該框架的核心在於引入了分子片段的概念，並在原子/分子片段兩個層次上分別建模短程和長程效應。

圖1：LSR-MP 框架示意圖

具體而言，LSR-MP 框架包含三個關鍵模組:短程模組、分割模組和長程模組。短程模組會採用 EGNNs，在原子的區域性鄰域內傳遞資訊，捕捉化學鍵、鍵角等幾何資訊。這一過程可以看作在由短程截斷半徑定義的子圖上進行訊息傳遞。分割模組則會使用 BRICS （Breaking of Retrosynthetically Interesting Chemical Substructures）演算法，基於化學啟發式規則，將分子非重疊地劃分為多個片段。每個片段的表示會透過求和其包含的原子表示而得到。本質上，分割模組是在原子和片段的粒度間建立對映。而長程模組主要負責捕捉片段間的相互作用。研究員們構建了一個原子-分子片段的二部圖，其中原子和分子片段之間的邊由一個較大的長程截斷半徑定義。在此二部圖上，研究員們還設計了方向性注意力機制，用於長程資訊傳遞。

BRICS 是一種基於規則的化學分割演算法，如圖2所示，常用於將大型分子劃分為合適的片段，以便進行量子化學計算或機器學習建模。其基本原理如下:

預定義化學鍵模式：BRICS 預定義了16種化學鍵模式，這些模式是由有機化學專家根據經驗總結的，代表了分子中相對容易斷裂的鍵。例如，碳-雜原子鍵，苄位 C-C 鍵等。
斷鍵：對於輸入的分子，BRICS 先識別其中符合預定義模式的化學鍵，將分子按照這些鍵斷開，來得到一系列片段。
片段修飾：斷鍵後的片段通常帶有懸掛鍵，BRICS 會在斷鍵處加上合適的原子（通常是氫原子），以使得片段在化學上穩定。
片段過濾：BRICS 還會對生成的片段進行後處理，包括去除過小的片段（如單個原子）、去除重複片段等，以保證片段集合的合理性。

與其他基於圖的分割演算法相比，BRICS 具有顯著優勢。它不僅能夠利用化學家經驗中預定義的化學鍵模式，生成合理且符合化學直覺的片段，而且能夠在分割過程中最大程度地保留分子的化學特徵，使得生成的片段在後續建模中更具代表性。此外，BRICS 作為一種基於規則的演算法，其計算效率較高，特別適合處理大規模資料集。

BRICS 演算法在片段生成領域被廣泛使用，並在多個分子模擬任務中取得了不錯的效果。在 LSR-MP 框架中，BRICS 扮演著重要的角色，負責在原子-片段兩個粒度之間建立聯絡，為長短程資訊傳遞提供基礎。

圖2：BRICS 框架示意圖

透過短程模組處理區域性相互作用，並用長程模組補足長程相互作用，LSR-MP 框架在處理大型分子體系時，平衡了計算效率和物理準確性。同時，分割模組充當了原子/分子片段粒度之間的橋樑，使該框架能夠靈活處理不同的尺度體系。

根據研究，LSR-MP 框架在數學上還具有等變性，這意味著它能很好地適應分子體系的平移、旋轉等對稱性，而這是精準分子模擬的關鍵。此外，透過引入基於片段的長程資訊傳遞，LSR-MP 框架能夠緩解圖神經網路中的過度擠壓問題，提高深層網路的學習能力。

高效能預測，跨模型效能提升

研究員們在 MD22 和 Chignolin 兩個廣泛使用的基準資料集上，將基於 LSR-MP 的 ViSNet-LSRM 模型與多個 SOTA 模型進行了系統比較，其實驗結果均提升顯著。

ViSNet-LSRM 在大多數分子的能量和力預測任務上，都取得了最優表現。以114個原子組成的 AT-AT-CG-CG 分子為例，ViSNet-LSRM 將力預測誤差降至0.11 kcal/mol·Å，較優於 Equiformer 等模型15-40%。在166個原子的 Chignolin 蛋白上，ViSNet-LSRM 的能量預測誤差比主流模型低了30%以上。

更重要的是，ViSNet-LSRM 在模型效率上同樣表現優秀。它的引數量僅為 Equiformer 的56%，達到了與之相當甚至更優的預測精度。推理速度上，ViSNet-LSRM 也比主流模型快了8~43倍。

此外，研究員們考察了 LSR-MP 框架的可拓展性（圖3）。實驗表明，當將 LSR-MP 整合到 PaiNN、Equiformer 等不同的 GNN 模型後，它們的預測誤差平均降低了30%以上。這證實了 LSR-MP 引入的物理資訊具有一定的“普世價值”，能適用於不同的模型。

圖3：LSR-MP 的可拓展性

最後，研究員們對一個核苷酸分子體系（AT-AT）進行了 20ps 的分子動力學模擬（單步步長為1fs），使用 ViSNet-LSRM 作為力場，並在恆能量集合（NVE）條件下進行。透過分析 AT-AT 分子的振動光譜（如圖4所示）可以看到，儘管與 DFT 結果相比存在微小的峰強度差異，但 MD22 資料集中的軌跡與 ViSNet-LSRM 模擬的軌跡展現出了相似的振動光譜。這一發現表明，LSR-MP 的模擬能夠準確再現分子在較長時間尺度上的實際振動模式。