編輯 | 蘿蔔皮
近年來，基於從頭算參考計算的機器學習力場 (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 模擬的穩定性以及更高溫度的外推行為有關。能夠證明，即使測試誤差估計平均相同，等變性也會降低誤差分佈的擴散。
因此，透過等變表示使用方向資訊在精神上與經典 ML 理論類似，其中對映到更高的維度可以產生更豐富的特徵空間，這些特徵空間更容易引數化。
未來的研究
在最近的一系列研究中，提出了旨在降低 SO(3) 卷積計算複雜度的方法。它們可以作為完整 SO(3) 卷積的替代品，而本文介紹的方法可以在訊息傳遞正規化中完全避免使用昂貴的 SO(3) 卷積。
這些結果都表明，等變相互作用的最佳化是一個尚未完全成熟的活躍研究領域，可能提供進一步改進的途徑。
雖然該團隊的工作，使得使用現代 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-ai-fundamental-physics.html