一個模型搞定元素周期表常見元素：中國團(tuán)隊(duì)打造分子模擬預(yù)訓(xùn)練模型，最高節(jié)省90%數(shù)據(jù)

白交 2022-08-27 13:34:43 來源：量子位

白交 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

分子模擬領(lǐng)域的預(yù)訓(xùn)練模型，來了！

DPA-1，中國團(tuán)隊(duì)深勢科技以及北京科學(xué)智能研究院等機(jī)構(gòu)打造，能覆蓋元素周期表大多數(shù)常見元素。

在各類數(shù)據(jù)集上的遷移學(xué)習(xí)結(jié)果表明，該模型能大幅降低新場景對數(shù)據(jù)的依賴，甚至在特定條件下能省去90%的數(shù)據(jù)。

用大模型的思路打開分子模擬

機(jī)器學(xué)習(xí)輔助下的原子間勢能面（PES）建模，與相應(yīng)的機(jī)器學(xué)習(xí)勢函數(shù)正在徹底改變分子模擬領(lǐng)域。PES是用于描述化學(xué)體系的一個基本量，通過它能得到大量原子間相互作用的信息。

過去傳統(tǒng)的分子模擬，主要是基于物理經(jīng)驗(yàn)得到的解析函數(shù)來描述PES，但始終無法完整描述出原子間相互作用。

目前，機(jī)器學(xué)習(xí)勢函數(shù)得益于第一性原理計(jì)算產(chǎn)生的數(shù)據(jù)，極大拓展了模型的應(yīng)用范圍，在材料科學(xué)、計(jì)算物理等領(lǐng)域應(yīng)用廣泛，并取得了較大成果。

不過仍面臨著現(xiàn)有模型遷移能力不足、缺乏通用大模型的局限性。當(dāng)面對一個復(fù)雜體系時，仍需獲取大量數(shù)據(jù)從頭開始訓(xùn)練模型，造成計(jì)算成本高昂。

基于這樣的背景，研究人員參考當(dāng)下在CV、NLP等領(lǐng)域中大模型的一種“預(yù)訓(xùn)練+少量數(shù)據(jù)微調(diào)”解決方案，提出了DPA-1，基于新注意力機(jī)制的深度勢能預(yù)訓(xùn)練模型。

類似于NLP的注意力機(jī)制，研究人員提出了門控注意力機(jī)制?（即圖中的紅色模塊）。

據(jù)介紹，該模型在原子局部環(huán)境矩陣上進(jìn)行了類比多體（大量粒子構(gòu)成的微觀系統(tǒng)）相互作用的信息交互，并以正則化的相對坐標(biāo)點(diǎn)乘作為角度信息，對獲得注意力權(quán)重進(jìn)行重新加權(quán)，以此來實(shí)現(xiàn)類似門控的機(jī)制。

此外，DPA-1還引入了對元素的編碼。不同元素共用同一套網(wǎng)絡(luò)參數(shù)，從而提升元素容量。

還將化學(xué)元素可視化

在遷移性測試中，研究人員有意將不同訓(xùn)練集劃分成多個子集，且每個子集的組分、構(gòu)型都有較大差異。

以AlMgCu合金數(shù)據(jù)集為例，則是分為了single、binary、ternary三個子集。

（single子集僅包含單質(zhì)/一元數(shù)據(jù)，binary僅有二元數(shù)據(jù)，即Al-Mg，Al-Cu，Mg-Cu；而ternary則是剩余的三元數(shù)據(jù)）

結(jié)果顯示，相較于DeepPot-SE，DPA-1的測試精度有較大的提升，甚至在特定條件下這種提升達(dá)到了一兩個數(shù)量級。

進(jìn)一步地，研究人員設(shè)計(jì)了遷移學(xué)習(xí)的方案。

簡言之，就是先在較大規(guī)模數(shù)據(jù)上進(jìn)行預(yù)訓(xùn)練，然后根據(jù)新的少量數(shù)據(jù)集修改最后一層能量偏差。

在AlMgCu合金數(shù)據(jù)集測試中，就將一元、二元子集上進(jìn)行預(yù)訓(xùn)練，然后在三元子集上測試。

結(jié)果顯示，對比DeepPot-SE，DPA-1可節(jié)省約90%的三元數(shù)據(jù)。

在僅有少量三元數(shù)據(jù)測試下，也能達(dá)到較高的精度。

接著他們在包含56種元素的大型數(shù)據(jù)集OC2M上預(yù)訓(xùn)練，并將其遷移到毫不相關(guān)的HEA和AlCu數(shù)據(jù)集上，結(jié)果都顯示出能成功的應(yīng)用。

其他類似模型GemNet（分子的通用方向圖神經(jīng)網(wǎng)絡(luò)）在同樣OC2M數(shù)據(jù)集上預(yù)訓(xùn)練，產(chǎn)生的模型有數(shù)十億的參數(shù)，訓(xùn)練時間需要數(shù)以千計(jì)的GPU hours。

而DPA-1只需要不到200個GPU hours 來訓(xùn)練不到一百萬個參數(shù)，并取得了不錯的結(jié)果（DPA-1和Gemnet-OC[34]的能量MAE為0.681 vs 0.286 eV）。

除此之外，他們還將元素可視化——模型中學(xué)習(xí)到元素編碼進(jìn)行了PCA降維并可視化。

可以看到，所有元素都呈螺旋狀分布，同周期元素沿著螺旋下降，同族元素則垂直螺旋方，恰好對應(yīng)的是元素周期表的位置。

據(jù)研究人員介紹，本次研究證明了“預(yù)訓(xùn)練+少量任務(wù)微調(diào)”流程的可行性。

接下來，他們將致力于勢函數(shù)自動化生產(chǎn)、自動化測試，也將關(guān)注像多任務(wù)訓(xùn)練、無監(jiān)督學(xué)習(xí)、模型壓縮、蒸餾等操作，方便用戶一鍵生成下游任務(wù)所需的勢能函數(shù)。

感興趣的旁友，可戳下方論文鏈接了解更多~

論文鏈接：
https://arxiv.org/abs/2208.08236
GitHub鏈接：
https://github.com/deepmodeling/deepmd-kit
參考鏈接：
https://mp.weixin.qq.com/s/cbTgntJmuntskQmTWjAINA