面向晶體結構預測的深度學習方法①

劉志威,王宗國,郭佳龍,王彥棡

1(中國科學院 計算機網絡信息中心,北京 100190)

2(中國科學院大學,北京 100049)

晶體結構搜索主要是產生新結構,基于能量最低原理篩選優質結構的過程.當前確定晶體結構能量有兩種主要計算方法,分子動力學[1]和密度泛函理論方法[2].前者代表軟件如Lammps 等,后者代表軟件如VASP[3,4]和CASTEP[5]等.目前,國內外已經報道了一些通過不同結構搜索算法,結合預測結構能量的分子動力學或密度泛函理論等方法的晶體結構預測軟件.如基于隨機搜索算法結合密度泛函理論的AIRSS[6]軟件,基于粒子群搜索算法結合密度泛函理論的CALYPSO[7],基于雜化算法結合密度泛函理論的XtalOpt[8]軟件,基于遺傳搜索算法結合分子動力學或密度泛函理論的USPEX[9]軟件,以及基于遺傳算法同時結合了密度泛函理論和分子動力學的自適應遺傳算法(Adaptive Genetic Algorithm,AGA)晶體結構預測軟件[10].這些軟件都是基于第一性原理或者分子動力學方法獲取結構基態能量,基于分子動力學方法,計算速度快,但是精度低.而基于密度泛函理論方法計算精度高,但是計算時間長,尤其是隨著晶體結構復雜度的上升,其計算時間呈指數增長.尋找一種既可以保證計算精度,又可以降低計算成本的材料研究方法成為迫切需要解決的問題.

1 相關工作

隨著高性能計算和深度學習的發展,機器學習在材料科學中越來越多的被使用,可以準確預測材料的性能,不斷縮小材料科學與計算科學之間的差距.哥倫比亞大學的Nongnuch Artrith 研究組利用神經網絡訓練原子勢函數,用于快速預測結構能量,并開發了原子能網絡軟件包(Atomic Energy Network,AENet)[11].普林斯頓大學Weinan 課題組利用深度神經網絡訓練用于分子動力學計算的多體勢函數[12-14],所訓練的相互作用勢可以直接用于Lammps 軟件使用,目前該方法已經封裝為軟件包DeePMD-kit.中國科學院計算機網絡信息中心劉學源等提出了用于結構預測的機器學習框架[15,16].德克薩斯大學奧斯汀分校的You 等提出了多種機器學習算法來預測透明導體的生成焓和能帶間隙,采用的機器學習算法如最小二乘法、逐步回歸法,還有基于樹的模型(隨機森林模型和梯度提升模型)[17].日本東京大學的Ryo Yoshida 課題組為了促進遷移學習的應用,開發了一個XenonPy.MDL的預訓練模型庫,這些預訓練模型捕獲了材料的公共特性.克服了材料數據集數量有限的問題,可以通過預先訓練好的模型,利用有限的訓練數據來預測相關材料目標屬性[18].西澳大利亞大學的James P.Doherty 等課題組使用一種快速監督機器學習算法代替第一性原理計算預測材料彈性模量,快速篩選新能量,這種算法結合了線性回歸、神經網絡以及隨機森林回歸[19].南京郵電大學的李斌課題組等構建了深度神經網絡模型,通過學習開放量子材料數據(OQMD)中的生成焓參數訓練模型,并運用模型預測未知材料的生成焓進行預測,達到量子力學軟件的精度[20].Lu 等通過訓練神經網絡模型,從儀器壓痕中提取金屬和合金的彈性特征,作者建立了單縮進、雙縮進和多重縮進求解逆問題的算法,并提出幾種multifidelity 方法求解逆壓痕問題[21].

綜合考慮計算的準確性和時效性,借助深度學習框架,基于密度泛函理論計算數據集,在結構預測過程中利用深度神經網絡構建材料結構與能量之間的構效關系模型,加快結構能量的預測,并將該方法嵌入結構預測軟件使用.

2 預測晶體結構能量的學習模型

2.1 神經網絡和訓練

神經網絡通常包含一個輸入層,一個輸出層以及若干隱藏層,輸入層只接收輸入信息,通常為數據樣本的特征向量,經過隱藏分層和輸出層的處理,由輸出層輸出結果.隱藏層的層數和節點數可以根據實際情況進行調整.神經網絡能夠通過學習得到輸入與輸出的映射關系,整個訓練過程就是不斷更新權重以及偏置值,使得模型預測越來越準確.通常每個層之間可以加入激活函數引入非線性因素.圖1為典型的單隱藏層神經網絡.

前向傳播過程:

其中,σ為激活函數,可以看出每個神經元的值由上一層的神經元的值與連接權重、偏置值以及激活函數共同決定.

反向傳播算法(Back-Propagation,BP)[22,23]是神經網絡中更新參數最經典的方法,其代碼邏輯如算法1.

?

2.2 神經網絡模型的主要技術

2.2.1 激活函數

在人工神經網絡中,激活函數用于表示上層神經元輸出與下層輸入的非線性映射函數,其主要作用就是增強神經網絡的非線性建模能力.如無特殊說明激活函數一般是非線性函數.如果不用激活函數,每一層的輸出都是上一層輸入的線性函數,無論神經網絡有多少層,最終的輸出都是初始輸入的線性變換,無法擬合現實中很多非線性問題.引入激活函數后,使得神經網絡可以逼近任何非線性函數,理論上可以擬合任意分布的數據.常用的激活函數有Sigmoid 函數[24]、tanh 函數、Rectified Linear Unit (ReLU)函數[25,26]、ELU 函數[27]、和Parametric Rectified Linear Unit(PReLU) 函數[28].圖2展示了這幾種激活函數的圖像.

2.2.2 優化算法

權值更新算法是深度學習模型訓練中相當重要的一個模塊,其決定參數更新的方向與速度,好的權值更新算法和適當的優化參數能使模型更快更準地收斂.PyTorch 包含一系列權值更新算法,其接口通用性好,也方便集成更加復雜的算法.PyTorch 內集成SGD[29]、Adadelta[30]、Adagrad[31]、Root Mean Square Prop(RMSprop)[32],Adaptive Moment Estimation (Adam)[33]、Averaged Stochastic Gradient Descent (ASGD)[34]、Rprop[35]、LBFGS[36]等十余種權值更新算法.

學習率和BATCH_Size是深度學習中重要的超參數.若學習率設置過小,則收斂非常慢,若設置過大,則梯度會在最小值附近來回震蕩.BATCH_Size的選擇決定了梯度下降的方向,BATCH_Size 選擇過小,隨機性會加大,多數情況難以收斂.所以在實踐中,人們會選擇較大的BATCH_Size,較大的BATCH_Size 可以減小訓練時間,提高內存利用率,梯度下降方向較為準確,不容易引起震蕩;但是,盲目增大BATCH_Size 也會導致如內存不夠、模型的泛化能力下降[37]、收斂到sharp minimum[38]、迭代的次數增大等一系列問題.

3 用于晶體結構預測的深度學習模型框架

3.1 模型訓練總體框架設計

3.1.1 主流程

晶體結構預測主流程如圖3所示,分為3 個主要模塊,遺傳算法(Genetic Algorithm,GA)、密度泛函理論(Density Functional Theory,DFT)計算和模型訓練(Train),3 部分并行運行.

圖3 主流程圖

(1)GA:用于結構搜索.并行運行7 個晶系(立方晶系、四方晶系、方晶系、六方晶系、正交晶系、單斜晶系、三斜晶系),利用神經網絡勢(或經驗模型)預測能量,推薦低能結構(每個晶系默認推薦16 個,每次循環共112 個結構)開展DFT 計算;

(2)DFT 計算:并行計算GA 部分推薦結構的能量等信息.剔除冗余結構,將計算的結構和能量等計算結果存入數據庫,并加入Train 模塊的訓練集;

(3)Train 模塊:更新訓練集后,訓練新模型,更新GA 部分里的神經網絡勢模型.

模型收斂性通過訓練集的大小和模型誤差兩個標準決定.首先,確定訓練集大小.若訓練集太小,神經網絡易過擬合,訓練集過多時間花費大,而且易造成欠擬合,經過多次實驗及效果,采用112×20 作為訓練集的最小量,用戶可以根據實際情況進行調整;其次是模型誤差,通過實驗發現,當模型驗證誤差降到0.08 左右即可很好的收斂,用戶也可以根據實際情況進行調整.當模型收斂后,不再進行DFT 計算及模型更新,只進行GA 搜索,收斂的模型存入到數據庫中.當GA 搜索迭代完成時(默認循環總次數為20),推薦最終的候選結構加入數據庫.

3.1.2 GA 模塊詳細流程

軟件的初始設置參數:GA的搜索迭代次數、結構內原子類型、各原子結構數量、原子體積等.此外,還需要準備描述各原子環境的指紋特征文件.

本文采用的遺傳算法,能量作為適應度值,即能量越低,結構越穩定.在每一輪迭代中,按照晶體結構的7 個晶系并行搜索,既保證了搜索的全面性也提高了搜索效率.具體過程是先根據能量大小進行排序,選擇種群中前三分之一的低能結構[39](種群大小默認為64),對坐標信息進行交叉變異操作,產生下一代結構(默認產生下一代結構數量為16).在迭代過程中,為了防止陷入局部最優解,一方面設置一定的變異概率(默認為0.05)基因突變,突破搜索瓶頸;另一方面,我們采用多組不同的隨機種子,即產生多組不同的初始種群,增加搜索多樣性.

需要注意的是,在沒有初始模型的情況下,通過經驗模型迅速構建初始樣本空間.一旦訓練出神經網絡模型,便采用神經網絡勢預測;參數i為GA 迭代過程中當前迭代輪數;在沒有訓練出收斂的模型時,GA 內搜索總次數可以設置為較小值(默認為30),可以較快的訓練出收斂模型.模型收斂后增加GA 內搜索總次數(默認為1500).最后,每個晶系推薦出16 個低能結構共112 個結構.具體流程如圖4所示.

圖4 GA 模塊流程圖(以立方晶系為例,其余晶系過程相同)

3.2 模型訓練流程實現及優化

如圖5所示,我們使用4 個神經網絡層,輸入層神經元數量根據提取的原子特征數量靈活選取,2 個隱藏層各為10 個神經元.輸入為能夠滿足原子平移性、對稱性、置換性的原子局域環境描述子,這些描述子可以轉化為神經網絡的輸入特征向量,目前,提供了切比雪夫多項式[40]及Behler和Parrinello 提出的對稱函數[41,42]兩種描述原子局域環境的描述方法.輸出為擬合的原子能量值.隱藏層和輸出層的神經元都有偏置參數,由于更新過程與權重類似,這里不再展開.

圖5 神經網絡架構

(1)前向傳播過程

隱藏層1 輸入:

式中,Wih1為輸入層i節點到隱藏層h1節點的權重系數,xi為原始結構特征輸入.

隱藏層1 輸出:

式中,σ1為激活函數.

隱藏層2 輸入:

式中,Wh1h2為隱藏層1的h1節點到隱藏層2的h2節點的權重系數.

隱藏層2 輸出:

式中,σ2為激活函數.

輸出層:

式中,Wh2k為隱藏層2的h2節點到輸出層k的權重系數.輸出原子的能量值.訓練過程中,我們采用均方根誤差(RMSE)作為代價函數.

(2)反向傳播:以第l個樣本為例,η為學習率

隱藏層2 到輸出層的權重更新:

隱藏層1 到隱藏層2的權重更新:

輸入層到隱藏層1的權重更新:

這樣整個神經網絡完成一次參數更新.

需要注意的是,由于我們將總能量拆分為單個原子的能量,最后進行累加,所以與傳統的神經網絡訓練過程不同,每個原子的能量并不是已知的,因此權重參數不能直接擬合到樣本中.我們對每個結構中的原子能量累加后再進行誤差計算和反向傳播.

3.3 關鍵技術的選取

3.3.1 測試數據集

測試數據集中包含了原子數為6、9、12、18和24的TiO2晶胞,隨機生成Ti 原子與O 原子的配比共計15 種,包含6056 個晶體結構.每種組分包含的晶體結構數目見表1.

表1 測試數據集TiO2 結構信息

3.3.2 激活函數的選取

為測試激活函數對材料結構預測軟件中訓練模型的影響,采用目前深度學習常用的Adam 優化算法,設置學習率為0.0001,設置BATCH_Size為32 測試不同激活函數對模型訓練的影響.

從圖6可以看出,相比于Sigmoid 函數,其他激活函數收斂效果明顯較好,tanh 在開始階段收斂較快,最終的收斂效果幾個激活函數大致相似,綜合考慮,我們后續實驗均采用tanh 激活函數.

圖6 不同激活函數的損失變化趨勢圖

3.3.3 優化算法選取

學習率和BATCH 大小的選取也會影響優化算法的效果,我們綜合分析了學習率和BATCH 大小對8 種優化算法的影響,采取5 種BATCH_Size,0.000 01-0.1 之間選取5 種學習率進行測試.測試過程中隨機選取了5 組隨機種子,圖7中結果為5 組數據的平均值.

圖7中Adadelta、Adam、Adamax、ASGD、LBFGS、RMSprop、Rprop、SGD 算法分別用灰色、黑色、藍色、青綠色、洋紅色、黃色、紅色、綠色表示;BATCH_Size的值1、8、16、32和Full 分別由圓點、實心方形、實心五角、倒三角、十字表示.

圖7 優化算法、BATCH_Size、學習率大小對模型性能的影響

可以看出在BATCH_Size為1和8 時,多數優化算法表現較好,同時也可以看出盲目增大BATCH_Size,模型的泛化能力會降低,并且可能陷入sharp minimum.Adadelta 適合較大的學習率,而RMSprop、Rprop、Adam、Adamax 適合選取較小的學習率.

學習率和BATCH_Size 搭配最優情況下,Adadelta、Adam、Adamax、RMSprop的損失值均能達到0.043左右,均能收斂到理想的模型.

3.4 訓練集優化

為剔除訓練集中的冗余結構,在模型訓練過程中,并不是將每次循環DFT 計算的所有結構都加入到訓練集中,通過兩步剔除冗余結構.首先,考慮到每次循環的GA 都是從頭搜索,所以不同循環輪次之間可能會產生相同的結構,對模型訓練改進意義不大.故剔除每個循環倫次的重復冗余結構,不再加入訓練集.然后,為提高模型訓練效率,設定一個最小誤差項,ΔE=min_loss*n_atom,其中n_atom為原子數.若模型預測的結構總能量與DFT 計算的總能量的差值低于此值,則認為此結構對模型訓練改進意義不大,該結構不再加入訓練集.

此外,為防止模型出現過擬合的情況,在模型更新中我們采用了機器學習中比較典型的早停法[23].

4 軟件性能分析

4.1 主要參數設置及搜索策略

為了測試所提出的模型的性能以及實驗參數的設置,我們以Si 單晶、TiO2和CaTiO3三種常見的晶體結構的搜索為例,對本文所提出的模型進行測試,并與實驗結果進行對比分析.

首先,確定實驗參數.Si 單晶和TiO2化合物的原子局域環境的展開式采用Behler和Parrinello 提出的不變對稱函數表示,搜索激活函數選用雙曲正切函數(tanh),優化算法為Adam,學習率設置為0.0001,BATCH_Size 設置為1.CaTiO3的原子局域環境的展開式采用基于切比雪夫多項式表示,激活函數選用雙曲正切函數(tanh),優化算法為Adamax,學習率設置為0.001,BATCH_Size 設置為8.然后,開展模型訓練和結構搜索.初始結構根據七大晶系分別搜索,七大晶系的搜索采用并行計算,每種晶系每個循環輪次推薦16 個低能晶體結構用于開展DFT 計算,然后加入訓練集去訓練更優的模型,直至模型收斂.將收斂的模型用于結構搜索.每個晶系默認按20 個隨機種子搜索結構,每個種子推薦出16 個候選結構,最后每個晶系按模型預測能量排序推薦出10 個低能結構,然后進行DFT 計算,根據計算能量大小推薦晶體結構.

4.2 實驗結果分析

采用遺傳算法搜索晶體結構,基于深度學習訓練模型預測晶體結構的能量值.通過該方法預測的結構晶格常數信息和對應的能量值與DFT 計算結果見表2,E為DFT 計算總能量.表2可以看出,3 種結構的晶格常數的最大誤差分別為:0、6.88%和0.在滿足均方根誤差小于0.08 eV/atom的模型下,預測的3 種穩定結構總能量與實驗測定穩定結構的總能量差值分別為:0.34 eV、1.41 eV 以及0.21 eV.預測晶體結構與實驗結構示意圖見圖8,結構中原子坐標見附錄A.

表2 3 種晶體結構預測結果與實驗結果對比(單位:eV)

從圖8可以看出,基于深度學習框架實現的晶體結構預測軟件,對于單胞中原子數目較少的Si 單晶和CaTiO3結構的能量預測都具有很好的效果,隨著元素種類的增加,模型誤差增大,因此對于單胞中原子數較多的原子還需要增大訓練集的多樣性.另外,通過圖8也可以看出,用于結構搜索的遺傳算法,隨著元素種類的增加,結構的搜索效果逐漸降低,CaTiO3結構的成鍵圖8(f)以及TiO2結構的原子點位圖8(d)與實驗結構有所偏移.因此,在結構搜索中需要增加種子數和迭代次數來修正.

圖8 預測晶體結構與實驗觀測結構

相比于基于密度泛函理論的第一性原理方法,神經網絡模型預測晶體結構能量節省了更多的計算資源和時間.為保證結構搜索的完整性,需要搜索的晶體結構數量巨大,如果所有結構都采用第一性原理方法計算結構能量幾乎是無法完成的,我們提出的面向材料結構預測的深度學習方法,只需對要少量有針對性的結構開展第一性原理計算,作為神經網絡訓練的數據,即可訓練出精度較高的能量預測模型,從而實現其他結構能量的快速預測.此外,隨著晶體結構復雜度的增加,第一性原理計算的時間成本呈指數增加,而我們提出的深度學習方法受結構復雜度影響較小,同樣只需要少量的第一性原理計算,構建高精度的神經網絡模型預測晶體結構能量.通過該方法在保證計算精度的前提下,既提高了計算效率,又克服了部分復雜體系無法開展大量計算的困難.

5 結論

本研究針對批量材料結構能量預測效率問題實現了一種用于材料結構預測的基深度學習方法,通過深入的分析和研究,對深度學習模型所需的訓練策略、測試集數據篩選方法、優化算法及其對應的學習率等參數進行了優化.并將深度學習框架應用到材料結構的預測中,預測的晶體結構參數與實驗發現晶體結構參數相吻合.該方法可以有效地提高材料結構預測效率,加快新材料的研發.

附錄A

(1)推薦結構原子位置與實驗制備結構對比(表A.1-表A.3).

表A.1 TiO2 結構原子位置

表A.2 CaTiO3 結構原子位置對比

表A.3 Si 原子位置對比

(2)預測不同晶系對應能量最低的晶體結構(表A.4).

表A.4 預測不同晶系對應能量最低的晶體結構