中國推出全新預訓練深度勢能模型DPA-1,對分子模擬產生深遠影響



「DPA-1 模型的提出證明了基於大模型實現『預訓練+少量數據微調』流程的可行性。這是是勢能函數生產的新範式,也是未來一系列工作的起點。」對於近期研究成果所具有的重要意義,一支中國團隊這樣表示。


基於機器學習的預訓練深度勢能模型正在分子模擬領域發揮著日益重要的作用,與此同時,由於現有的模型遷移能力有限,訓練成本較高,對訓練數據的依賴性也很強,所以其在實際應用中的表現仍然不盡人意。雖然相關研究人員為了解決這些問題也做出了很多探索和實踐,但成效尚不明顯。


為了克服上述因素造成的不良影響,解決在分子模擬背景下,面對新的複雜體系仍需生成大量數據來從頭訓練模型的關鍵難題,由深勢科技(DP Technology)和北京科學智能研究院(AI for Science Institute,Beijing)研究員及合作者組成的中國團隊基於新的門控注意力機制(Gated Attention Machanism),推出了通用性較強,能容納元素周期表中大多數元素的模型 DPA-1。


近日,相關論文以《基於注意力機制的分子模擬預訓練深度勢能模型》(DPA-1:Pretraining of Attention-based Deep Potential Model for Molecular Simulation)為題在
arXiv 上預發表[1]。



(來源:arXiv)




DPA-1 模型是在 DP 系列模型基礎之上的全面升級,具有以下優勢。


首先,該模型利用與自然語言處理領域注意力機制比較相似的門控注意力機制,對原子之間的相互作用進行了充分的建模,這能使模型在現有的數據條件下學習更多隱含的原子交互信息,可有效提升模型在不同數據集之間的遷移能力和數據生成時的採樣效率。


其次,模型中包含了經過編碼後的元素,且不同元素用的是相同的網路參數,這有利於拓展模型中的元素容量。


同時,由於模型在擁有 56 種元素的大數據集上開展了預訓練,並在多個下游任務上完成了遷移學習,所以能夠在保證預測精度的前提下,大大降低訓練成本和訓練數據量。


此外,該模型還擁有超高的推理效率,可執行大規模的分子動力學模擬。



▲圖 | DPA-1 模型示意圖(來源:arXiv)



為了切實有效地避免傳統模型帶有的局限性,開發人員開展了幾項有針對性的實驗。


開發人員先將不同訓練集以多子集形式進行劃分,然後在訓練一部分子集的同時去測試另一部分子集。需要說明的是,這裡每個子集之間的構象和組分都不相同,比如,在 AlMgCu 數據集上,single 子集中只有單質數據,binary 子集中只有二元數據,ternary 子集中只有三元數據。


最後,開發人員分別對 DPA-1 和 DeepPot-SE 這兩個模型在 AlMgCu 合金、固態電解質(SSE,solid state electrolyte)和高熵合金(HEA,High-entropy alloys)這三類數據集上的表現進行了測試。結果顯示,與 DeepPot-SE 相比,DPA-1 的測試精度能達到一兩個數量級的提升,這充分說明了後者擁有強大的遷移能力。



▲圖 | 在不同訓練集上測試得到的結果(來源:arXiv)



在「預訓練+少量數據微調」的模型生產範式下,開發人員給 DPA-1 規划了一套遷移學習方案。先在大規模數據上開展模型預訓練工作,而後借新數據集的統計結果修改最後一層的能量偏差,並將其作為新任務的訓練起點。


比如,先在 AlMgCu 數據集中的一、二元數據上執行預訓練,並在三元數據上完成測試。緊接著,執行 OC2M 數據集上的預訓練工作,再分別遷移至 HEA 和 AlCu 數據集上。結果顯示,DPA-1 不僅能在只有三元數據的場景下實現較高精度,還能有效減輕對下游訓練數據的依賴。



▲圖 | 在不同數據集上,DPA-1 和 DeepPot-SE 的學習曲線圖(來源:arXiv)



開發人員還將 DPA-1 中已被編碼的元素參數進行了 PCA 降維和可視化表現。結果表明,在隱空間中所有的元素都呈螺旋狀分布的態勢,並且同周期的元素沿螺旋下降趨勢分布,同族的元素垂直於螺旋分布,這種分布態勢與其在元素周期表中的位置巧妙對應,能夠很好地證明模型的可解釋性。



▲圖 | PCA 降維和可視化表現圖(來源:arXiv)



目前,該團隊已在其科學計算雲平台 Bohrium 上完成了 DPA-1 的開源工作,DPA-1 關於訓練和分子動力學模擬功能的開源也已在 DeepModeling 開源社區的 DeePMD-kit 項目下實現。


該團隊表示:「未來,我們將繼續致力於勢能函數的自動化生產和自動化測試研究,仍會繼續關注諸如多任務訓練、無監督學習、模型壓縮和蒸餾等方面的操作。此外,更大更全的資料庫、下游任務和 dflow 工作流框架的結合也是著重發展的方向。」


參考資料:
1.Duo, Z., Hang, B.et al. DPA-1:Pretraining of Attention-based Deep Potential Model for Molecular Simulation.
arXiv (2022). https://arxiv.org/abs/2208.08236