2.5天完成1年的MD計算？ DeepMind團隊基於歐幾裡得Transformer的新計算方法

編輯 | 蘿蔔皮

近年來，基於從頭算參考計算的機器學習力場 (MLFF) 的開發取得了巨大進展。雖然實現了較低的測試誤差，但由於擔心在較長的模擬時間範圍內會出現不穩定性，MLFF 在分子動力學 (MD) 模擬中的可靠性正面臨越來越多的審查。

研究表明，對累積不準確性的穩健性與 MLFF 中使用等變表示之間存在潛在聯繫，但與這些表示相關的計算成本可能會在實踐中限制這種優勢。

為了解決這個問題，Google DeepMind、柏林工業大學（TU Berlin）的研究人員提出了一種名為SO3krates 的transformer 架構，它將稀疏等變表示（歐幾里德變數）與分離不變和等改變訊息的自註意力機制結合，因此無需昂貴的張量積。

SO3krates 實現了精確度、穩定性和速度的獨特組合，能夠對長時間和系統尺度上的物質量子特性進行深入分析。

研究以「A Euclidean transformer for fast and stable machine learned force fields」為題，於 2024 年 8 月 6 日發佈在《Nature Communications》。

背景與挑戰性

分子動力學（MD）模擬透過長時間尺度的模擬，可以揭示系統從微觀相互作用到宏觀性質的演變，其預測精度取決於驅動模擬的原子間力的精確度。傳統上，這些力來自近似的力場（FF）或計算複雜的從頭計算電子結構方法。

近年來，機器學習（ML）位能模型透過利用分子系統的統計依賴性，提供了更靈活的預測手段。

然而，研究表明，ML模型在基準資料集上的測試誤差與長時間尺度MD模擬中的表現相關性較弱。

為改善外推性能，訊息傳遞神經網路（MPNNs）等複雜架構被開發，特別是等變MPNNs，透過引入張量積捕捉原子間的方向訊息，提高了數據的可轉移性。

在 SO(3) 等變架構中，卷積是在球諧函數基礎上對 SO(3) 旋轉群進行的。透過固定架構中球諧函數的最大次數，可以避免相關函數空間的指數成長。

科學家已證明最大階數與準確度、數據效率密切相關，並與 MD 模擬中模型的可靠性相關。然而，SO(3) 卷積的規模為，與不變模型相比，這可以將每個構象的預測時間增加多達兩個數量級。

這導致了一種必須在準確性、穩定性和速度之間做出妥協的情況，也可能會帶來重大的實際問題。必須先解決這些問題，這些模型才能在高通量或廣泛的探索任務中發揮作用。

性能強勁的新方法

Google DeepMind、柏林工業大學的研究團隊以此為動機，提出了一種歐幾里德自註意力（Euclidean self-attention）機制，用原子鄰域相對方向的過濾器取代SO(3) 卷積，因此無需昂貴的張量積即可表示原子交互作用；此方法稱為SO3krates。

圖示：SO3krates架構和構建塊。 （資料來源：論文）

此解決方案建立在神經網路架構設計和幾何深度學習領域的最新進展之上。 SO3krates 使用稀疏表示來表示分子幾何形狀，並將所有捲積響應的投影限制在等變基函數中最相關的不變分量上。

圖示：學習不變量。 （資料來源：論文）

由於球諧函數的正交性，這種投影對應於乘積張量的跡，可以用線性縮放內積來表示。這可以高效地擴展到高階等變表示，而不會犧牲計算速度和記憶體成本。

力預測是根據所得的不變能量模型的梯度得出的，該模型代表自然等變的分段線性化。在整個過程中，自註意力機制用於分離模型中的不變和等變基礎元素。

該團隊比較了 SO3krates 模型的與目前最先進的 ML 模型的穩定性和速度，發現該解決方案克服了當前等變 MLFF 的局限性，同時又不損害它們的優勢。

研究人員提出的數學公式可以實現高效的等變架構，從而實現可靠穩定的 MD 模擬；與具有相當穩定性和準確性的等變 MPNN 相比，其速度可提高約 30 倍。

為了證明這一點，研究人員僅在幾個小時內就對超分子結構運行了精確的納秒級MD 模擬，這使他們能夠計算出從具有42 個原子的小肽到具有370 個原子的奈米結構的結構的收斂速度自相關函數的傅立葉變換。

圖示：結果概述。 （資料來源：論文）

研究人員進一步應用該模型，透過使用最小值跳躍演算法研究 10k 最小值，探索二十二碳六烯酸 (DHA) 和 Ac-Ala3-NHMe 的 PES 拓撲結構。

這樣的研究需要大約 30M FF 評估，這些評估在幾百 K 到 1200 K 之間的溫度下進行。使用 DFT 方法，這種分析需要一年以上的計算時間。現有的具有類似預測精度的等變 MLFF 需要運行一個多月才能完成此類分析。

相比之下，團隊僅用 2.5 天就能完成模擬，從而有可能在實際時間尺度上探索數十萬個 PES 最小值。

此外，SO3krates 能夠偵測訓練資料中未包含的物理有效最小構象。外推到 PES 未知部分的能力對於將 MLFF 擴展到大型結構至關重要，因為可用的從頭算參考資料只能涵蓋構象豐富的結構的子區域。

該團隊還研究了在網路架構中禁用等方差屬性的影響，以更深入地了解其對模型特性及其在 MD 模擬中的可靠性的影響。

研究人員發現，等變性與所得 MD 模擬的穩定性以及更高溫度的外推行為有關。能夠證明，即使測試誤差估計平均相同，等變性也會降低誤差分佈的擴散。

圖示：計算效率和 MD 穩定性。 （資料來源：論文）

因此，透過等變表示使用方向資訊在精神上與經典 ML 理論類似，其中映射到更高的維度可以產生更豐富的特徵空間，這些特徵空間更容易參數化。

未來的研究

在最近的一系列研究中，提出了旨在降低 SO(3) 卷積計算複雜度的方法。它們可以作為完整 SO(3) 卷積的替代品，而本文介紹的方法可以在訊息傳遞範式中完全避免使用昂貴的 SO(3) 卷積。

這些結果都表明，等變相互作用的最佳化是一個尚未完全成熟的活躍研究領域，可能提供進一步改進的途徑。

圖示：穩定、長時間尺度的 MD 模擬和對更大生物分子的外推。 （資料來源：論文）

雖然團隊的工作，使得使用現代 MLFF 建模範式可以實現穩定的擴展模擬時間尺度，但未來仍需進行最佳化，以使 MLFF 的適用性更接近傳統的經典 FF。

目前，該方向出現了各種有潛力的途徑：在當前設計中，EV 僅根據二體相互作用進行定義。透過將原子簇擴展納入 MP 步驟，可以進一步提高準確性。同時，這可能有助於減少 MP 步驟的數量，從而降低模型的計算複雜度。

另一個尚未討論的問題是全局效應的適當處理。透過使用低秩近似、可訓練的 Ewald 求和或以物理啟發的方式學習長程校正。當需要外推到更大的系統時，後一種類型的方法尤其重要。

雖然等變模型可以提高局部交互作用的外推能力，但對於超出訓練資料中存在的長度尺度或超出模型有效截止值的交互作用，這並不適用。

由於上述方法依賴於局部屬性，例如部分電荷、電負性或赫什菲爾德體積，因此可以透過學習SO3krates 架構的不變特徵分支中的相應局部描述符，將它們無縫整合到該團隊的方法中。

因此，未來的工作將集中於將多體展開、全局效應和長程相互作用納入 EV 形式，並旨在進一步提高計算效率，最終以高精度跨越 MD 時間尺度。

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

相關內容：https://phys.org/news/2024-08-faster-coupling-fast ai-fundamental-physics.html

以上是2.5天完成1年的MD計算？ DeepMind團隊基於歐幾裡得Transformer的新計算方法的詳細內容。更多資訊請關注PHP中文網其他相關文章！