速度暴漲1000倍!擴散模型預測材料“煉金”過程,體系越復雜加速度越大
用擴散模型預測化學反應,速度直接暴漲1000倍!
原本需要用計算機硬算幾小時甚至一天,現在單個GPU用6秒鐘就能搞定。
這是MIT和康奈爾大學聯合搞出來的一項新研究,用擴散模型來預測化學反應中最關鍵的過渡態結構,不僅計算速度提升1000倍,結果也竟然意想不到得準確,相關研究工作發表在Nature Computational Science上:
其中,麻省理工學院的段辰儒博士是第一加通訊作者。此外,康奈爾大學博士生杜沅豈、麻省理工學院博士生賈皓鈞以及麻省理工學院Heather Kulik教授為該論文的共同作者,目前研究已經被MIT News報道。
要知道,預測反應中的過渡態結構遠非想象中簡單——
由于能量較高,它存在的時間往往只有飛秒級,即秒的千萬億分之一。
因此,目前還沒有一臺實驗設備能直接觀察它,只能通過量子化學計算的方式來預測,計算量少則幾小時,多則數天。
如今用擴散模型就能得到幾乎一樣的預測結果,這究竟是怎么做到的?
進一步地,用AI預測化學反應,又究竟可以被應用在哪些領域、起到哪些作用?
我們和論文的其中兩位作者,來自MIT的段辰儒和康奈爾大學的杜沅豈聊了聊,探討了一下這項研究具體的原理、背后潛在的應用方向以及關于AI for Science的一些思考。
如何用擴散模型預測化學反應?
首先,需要理解為啥過渡態是研究化學反應的關鍵。
自遠古的“煉金術”以來,化學一直是一門了解和控制物質之間相互作用的學科,化學反應又是其中非常核心的概念。
通常來說,一個化學反應由三個東西組成:反應物、生成物和過渡態結構。
反應物和生成物大家都很熟悉了(如氫氧生成水),但實際上,二者之間的過渡態才是解釋化學反應的關鍵——
作為化學反應過程中能量最高的狀態,過渡態既能用來理解化學反應的機理、也能估算反應速率和能量。
一言以蔽之,要想真正搞懂、設計、優化并調控一個化學反應,就必須從過渡態結構下手。
所以,過渡態結構究竟要如何研究,又為什么需要用AI來做預測?
研究過渡態結構,不能只用分子結構來分析,而要更進一步研究它的3D構象。
分子結構,指為了便于理解化學反應,會人為構造出單雙鍵這樣的概念,并用鍵合關系、原子種類和數量來表示化學反應。
但要想真正精確計算化學反應的速率和能量,就必須要研究原子在3D空間中的位置關系(用3D坐標表示),即3D構象。只有對比原子在3D空間中的位置關系,才能進行量化分析,從而預測反應發生的過程。
然而,相比反應物和生成物,過渡態存在的時間非常短,甚至只有飛秒量級。
受限于實驗設備精度,過渡態結構無法用肉眼直接觀察,此前只能通過量子力學方法——薛定諤方程來計算。
薛定諤方程,量子力學中的基本方程,描述粒子在某段時間內的狀態如何變化。
BUT,手搓薛定諤方程計算量巨大,一個苯環的化學反應可能都需要算上一周時間,尤其隨著體系增大(原子數量變多等),計算量更是呈現出指數級增長的趨勢。
雖然后來出現了密度泛函理論(DFT)等近似方法加速計算,但算起來還是很慢、有時候甚至因為收斂性的問題導致計算不出來,“成本很昂貴,出錯率還高”。
為此,來自MIT和康奈爾大學的研究人員,想到用AI的方法來直接預測過渡態結構,以節省計算量。
之所以選用擴散模型,據杜沅豈介紹,主要有三點原因:
其一,生成效果好。相比GAN,VAE等模型,擴散模型算是生成模型中效果較好的架構。
其二,更適合3D構象。相比分子結構是離散的數據,3D構象是一個連續的數據,且結構上涉及(平移、旋轉和置換)對稱性。
對于離散數據而言,基于機器學習+搜索的模型會比現有的(離散)擴散模型效果更好;但對于3D構象而言,無論是其對稱性還是用連續數據表示,反而更適合用擴散模型生成。
其三,擴散模型自身的靈活性,使得它可以在生成時,保留一部分原有設計、填補剩余的部分,也使得它對于設計新的化學反應更有幫助。
具體到架構設計上,作者們將等變GNN融入到擴散模型中,設計了一種名叫OA-ReactDiff的結構,其中等變GNN能很好地處理并保留數據的對稱性。
所以,OA-ReactDiff的效果究竟如何,或者說相比其他AI最大的亮點究竟是?
段辰儒表示,這項研究最核心的亮點主要有兩個。
一個是預測準確性上,OA-ReactDiff的準確性超過了之前的一系列AI模型。
作者們選用了Transition1x作為數據集,這個數據集包含10073個化學反應,每個化學反應分別包含反應物、生成物和經過量子化學計算的過渡態結構,整個體系不超過23個原子。
這其中的9000個化學反應用作模型訓練,1073個作為測試集,最終實現了0.183?(1?=0.1納米)的平均均方根偏差和0.076?的中位數均方根偏差。
相比于其他機器學習方法如PSI-based和NeuralNEB,OA-ReactDiff準確性都要更高。
與此同時,在沒有刻意訓練的情況下,OA-ReactDiff在更大的體系上也表現出了良好的預測能力。
另一個是預測速度上,OA-ReactDiff相比現有的計算方法,至少能快上1000倍。
這是因為,原本密度泛函理論算法的時間復雜度在O(N3)級別,但現在基于OA-ReactDiff的方法,直接將時間復雜度降低到O(N2)。
反映到現有數據集上,原本需要計算幾小時甚至一天的化學反應,現在用OA-ReactDiff,只需要6秒鐘就可以搞定。
進一步地,體系越大的化學反應,用擴散模型做預測的提速效果還會越好。
不過,這項研究究竟能被用在哪些地方呢?
我們也問了問兩位作者關于這項研究潛在的應用、以及對于“用AI搞科學研究”這件事本身的看法。
可用于燃料藥物輔助設計
對于研究本身,段辰儒認為主要有三個方面的作用。
最直觀的一方面,自然是催化劑的設計了。
這也是研究人員接下來想要繼續探索的方向,即繼續擴展模型本身,加入催化劑等條件,并用AI來預測過渡態。
而在研究催化劑的基礎上,又能進一步輔助燃料設計和新藥開發。
以燃料為例,雖然涉及的燃燒反應體系往往不大,但發生反應的過程卻極易被環境和條件影響。
因此,如果用AI來輔助預測化學反應的過程,或許能更快基于不同燃料發生反應的環境條件來預測效果。
這樣一來,對于能源行業甚至航天行業也會產生進一步的影響。
最后,就是模擬自然界化學反應,來幫助開發新的理論研究。
像是模擬地球生命早期可能發生的反應、或是探索行星氣體之間的相互作用等,通過用AI預測過渡態,也可能更快發現或驗證新的理論,探究生命起源的意義。
聽起來,AI似乎已經給化學研究帶來了不少新的突破。
結合這段時間AI在生物制藥(AlphaFold2)、以及材料學(GNoME)上取得的種種突破,對于生化環材這一系列理論學科而言,AI是否已經成為了新的技術主流路線?
更激進一點來說,對化學本身而言,未來AI是否會取代一部分現有的量子化學研究?
對此,兩位作者均持有相同的觀點,即AI(至少就監督模型而言),現階段對于這些學科的影響力依舊是一種工具,而并非一條全新的技術路徑。
這是因為,目前生化環材任何一個學科借助AI輔助研究的前提,依舊是已經有了一定的理論,而AI作為工具起到的是“錦上添花”的作用。
換言之,AI與這些學科的其他研究方法依舊是共存關系,并不會出現“替代”的情況。段辰儒對此做了一個有意思的比喻:
如果化學學科好比一個夾饃,機器學習方法和各種現有的化學方法,就是不同的食材。那夾饃里面可以夾肉,也可以夾其他的菜。
以近期大火的Google Deepmind研究GNoME為例,就是一個很好的“AI當工具人”的例子。
杜沅豈認為,如果從ML專業的角度去看待GNoME中涉及的AI方法,會發現它其實并不是最新的東西,像是其中涉及的幾何深度學習,已經是這幾年比較流行的研究方向。
但如果從材料學科的角度去看,這樣大規模、高通量的穩定材料搜索,肯定還是有相當的科學意義在其中:
AI for Science本身,其實更側重的是“Scientific Discovery”,它更大的意義在于科學發現。
而段辰儒認為,從應用的角度而言,無論是材料發現,還是預測化學反應,最終用AI去加速落地的宗旨依舊是“有用”二字:
正如這屆NeurIPS AI4Science的workshop主題一樣,“從理論到實踐“。我認為這是AI4Science后面十年的必經之路,也是AI4Science的中期使命。
像OA-ReactDiff接下來的目標是找到一個可以應用的路徑一樣,GNoME發現220萬穩定材料也只是應用的一個起點,接下來還需要很多的標準(如導熱、導電等),才能讓化學和材料真正在某一行業“發光發熱”。